文字通り優先度が付いたキューだ。通常のキューはFIFOだがこのキューは値が小さいものから順にpopできる。

Pythonではheapqが用意されているので、特段独自実装する必要がない。名前がヒープなのかキューなのかという感じだが、内部ではヒープになってるっぽい。使い方は以下の通りでheapifyでリストを優先度付きキューに変えられる。heappushで要素を入れ、heappopで要素を出す。

import heapq

list = [2,6,1]

heapq.heapify(list) 

heapq.heappush(list, 4)
heapq.heappush(list, 5)
heapq.heappush(list, 3)

// 1, 2, 3 と出力される。
print(heapq.heappop(list))
print(heapq.heappop(list))
print(heapq.heappop(list))

これによりリストを昇順に並び替えて先頭N個を何かしら計算してみたいなシチュエーションでO(logN)と高速に計算することができる。

一方で降順に並び替えてという場合は、リスト内の値を*-1してあげればよい。

例題: ABC141 D

D - Powerful Discount Tickets

要約するとN個の商品に対してM枚の割引券を適用した時に一番お得に買い物がしたいということだ。前回のエントリが貪欲法だったわけだが、今回も貪欲法を使うことになる。つまり一番高い商品に対してM回割引券を使えばよいのだ*1。割引券を使う対象は優先度付きキューで降順管理し、最後に負数を逆転すればいい。計算量はO(NlogN)。

僕は以下のように実装した。

import heapq

n,m=map(int, input().split())
a=list(map(int, input().split()))
minus_a=[-x for x in a]

heapq.heapify(minus_a)

for _ in range(m):
  num=heapq.heappop(minus_a)
  heapq.heappush(minus_a, (-1)*(-num//2))  # 負数の剰余演算を避けるため

print(-sum(minus_a))