優先度付きキュー(priority_queue)の勉強をしたので。
heap型を使えば良い。
heapはリストに比べて最小値の取り出しがはやく、新たな数の追加もはやい。
pythonにはheapqというmoduleがある。
heapqのheapには、基準となる数以外に別の数をtupleの形で持つことができる。
参考文献はこれ
8.5. heapq --- ヒープキューアルゴリズム — Python 3.7.0 ドキュメント。
今回解いた問題はcode festival2017のC問題です。
問題はこれ
C: Factory - CODE THANKS FESTIVAL 2017(Parallel) | AtCoder。
解説はこれ
https://img.atcoder.jp/code-thanks-festival-2017-open/editorial.pdf
解答はこれ。
import heapq N, K = map( int, input().split()) H = [tuple( map( int, input().split())) for _ in range(N)] ans = 0 heapq.heapify(H) #heap型にする for _ in range(K): M = heapq.heappop(H) #最小のものを取り出す a, b = M[0], M[1] ans += a heapq.heappush(H,(a+b,b)) #加算したものをheapに追加する print(ans)
おまけ
ずっと
H = heapq.heappush(H,(a+b,b))
ってやっててエラーが出続けていた。