優先度付きキューとは
優先度付きキュー（priority queue）とは、優先度にしたがって、優先度の高いものから順番に取り出すことができるものだ。
その代表例がヒープキュー（heapq）である。公式ドキュメントはこちら。（https://docs.python.org/ja/3/library/heapq.html）
特徴は大きく分けて2つある。
- 最小値、最大値を素早く取り出すことができる
- 値の挿入順は気にしなくても良い
主に使うメソッドは次のとおりだ。
-
heapq.heapify(heap, item)：優先度付きキューの作成 -
heapq.heappop(heap)：最小値の取り出し -
heapq.heapush(heap, item)：キューへの挿入
heapq.heappop(heap)、heapq.heapush(heap, item)の計算量はどちらも$O(logN)$なので、通常のリストでの操作よりも少ない計算量で実装できる。
# heapq
import heapq
numbers = [3, 5, 1, 6, 2, 4]
heapq.heapify(numbers) # 優先度付きキューの作成
print(numbers)
print(heapq.heappop(numbers)) # 最小値の取り出し
print(heapq.heappop(numbers))
print(numbers)
heapq.heappush(numbers, 1) # キューへの挿入
heapq.heappush(numbers, 2)
print(numbers)
print(heapq.heappop(numbers)) # 最小値の取り出し
print(heapq.heappop(numbers))
print(numbers)
>> [1, 2, 3, 6, 5, 4]
>> 1
>> 2
>> [3, 4, 5, 6]
>> [1, 3, 2, 6, 4, 5]
>> 1
>> 2
>> [3, 4, 5, 6]
例題1: ABC137 D - Summer Vacation
def main():
import heapq
n, m = map(int, input().split())
jobs = [[] for _ in range(m+1)]
for _ in range(n): # 各日のリストに各仕事の報酬を追加
a, b = map(int, input().split())
if a <= m:
jobs[a].append(b)
heap = []
heapq.heapify(heap)
total = 0
for day in range(1, m+1):
for i in jobs[day]:
heapq.heappush(heap, -i) # 最大値を取れるように-iでpush
if len(heap) > 0: # もし値があれば、
total += abs(heapq.heappop(heap)) # 最大値を取得
print(total)
if __name__ == '__main__':
main()
例題2: ABC141 D - Powerful Discount Tickets
def main():
import heapq
n, m = map(int, input().split())
alist = list(map(lambda x: int(x)*(-1), input().split()))
heapq.heapify(alist)
for _ in range(m):
temp = heapq.heappop(alist)*(-1)//2
heapq.heappush(alist, temp*(-1))
print(sum(map(lambda x: -x, alist)))
if __name__ == '__main__':
main()
例題3: ABC212 D - Querying Multiset
def main():
import heapq
q = int(input())
sack = []
heapq.heapify(sack)
total = 0
for i in range(q):
query = list(map(int, input().split()))
if query[0] == 1:
heapq.heappush(sack, query[1]-total)
elif query[0] == 2:
total += query[1]
else:
ans = heapq.heappop(sack)
print(ans+total)
if __name__ == '__main__':
main()
例題4: 第二回アルゴリズム検定 F - タスクの消化
def main():
import heapq
n = int(input())
tasks = [[] for _ in range(n+1)]
for _ in range(n):
a, b = map(int, input().split())
tasks[a].append(b)
heap = []
heapq.heapify(heap)
total = 0
for day in range(1, n+1):
for i in tasks[day]:
heapq.heappush(heap, -i) # 最大値の取得
if len(heap) > 0:
total += abs(heapq.heappop(heap)) # 最大値を加算
print(total)
if __name__ == '__main__':
main()
例題5: ABC173 D - Chat in a Circle
def main():
import heapq
n = int(input())
alist = list(map(int, input().split()))
alist.sort(reverse=True)
total = 0
heap = []
heapq.heappush(heap, alist[0]*(-1)) # 最初に到着した人
for i in range(1, n):
friendly = heapq.heappop(heap)*(-1) # 輪の中で最大値を取る
total += friendly # 心地よさに加算する
heapq.heappush(heap, alist[i]*(-1)) # 輪に次の人を入れる
heapq.heappush(heap, alist[i]*(-1))
print(total)
if __name__ == '__main__':
main()
例題6: ABC217 E - Sorting Queries
def main():
from collections import deque
import heapq
q = int(input())
que = deque()
heap = []
for _ in range(q):
query = list(map(int, input().split()))
if query[0] == 1:
que.append(query[1]) # queに値を追加
elif query[0] == 2:
if len(heap) > 0: # heapに値が存在したら
print(heapq.heappop(heap)) # heapの最小値を取り出す
else: # heapに値がなかったら
print(que.popleft()) # queから値を取り出す
else:
while que: # queの中身をすべてソートながらheapに挿入
heapq.heappush(heap, que.pop())
if __name__ == '__main__':
main()
例題7: ABC223 D - Restricted Permutation
heapqを用いたトポロジカルソート。
def main():
from heapq import heappush, heappop
n, m = map(int, input().split())
indeg = [0] * (n+1) # 入次数
out = [[] for _ in range(n+1)]
for _ in range(m):
a, b = map(int, input().split())
indeg[b] += 1
out[a].append(b)
heap = []
# 入次数が0の数字を探す
for i in range(1, n+1):
# 入次数が0ならば
if indeg[i] == 0:
heappush(heap, i)
ans = []
while heap: # すべての数字を見るまで続行
# 辞書順で最小
# つまり、数字が小さいものから取り出す
now = heappop(heap)
ans.append(now)
# 次の数字はどこへ向かうか
for to in out[now]:
indeg[to] -= 1
# 入次数が0の数字をheapに追加
if indeg[to] == 0:
heappush(heap, to)
if len(ans) == n:
print(*ans)
else:
print(-1)
if __name__ == '__main__':
main()
まとめ
heapq は実装こそ難しくないので、使い慣れると最大の武器となる。
問題に行き詰まったら、heapが使えないか確認する癖をつけておこう。