先日の ABC137 の D 問題を解く際に、 Python の heapq について軽く調べたのでメモを残しておきます。
そもそも heapq ってなんやねーん
先にざっくり概要を説明すると、 heapq というのは python のモジュールで、プライオリティキュー（優先度付きキュー）というデータ構造を実現するときに利用されます。
heapq は次の操作を何度も繰り返すようなときに強さを発揮します。
– 要素を追加するとき
– 追加した要素の中から最小値を取り出すとき
具体的には、要素の追加と削除を \(O(\log{N})\) で、最小値の取得を \(O(1)\) で行えます。
heap について
そもそも heap （ヒープ）とはグラフの木構造の 1 つで、これは heapq （プライオリティキュー、優先度付きキュー）を実現するときによく使われるものです。（要するに heap は heapq のベースみたいなもんです）
heap におけるデータを格納するときのルールとして、子要素の数字は親要素の数字よりも必ず大きいというものがあります（ヒープ条件）。これによって、木の一番上（つまり根の部分）に最小値がくるという特徴があります。
heap では常に一番上（根）に最小のデータがあるようにされるため、データの数に関係なく \(O(1)\) で最小値を取り出せることが特徴です。
heapq について
heapq というのは heap によって実現されるデータ構造で、データを自由に追加でき、取り出すときには最小値から順に選ばれるようになっていることが特徴です。
キュー（FIFO: 先入れ先出しのデータ構造）やスタック（LIFO: 後入れ先出しのデータ構造）のような、追加した順番で要素を保持するのとは異なり、 heapq では要素を最小値が先頭に来るようにソートして保持します。
heapq を使ってみる
まずは heapq モジュールを import します。
import heapqheapq を活用するためにまず heap 化されたリストを用意する必要がありますが、これには次の 2 つの方法があります。
hq = [] # 初期化されたリストを新規に作成
heapq.heapify(hq) # 既存のリストhqをheap化python の heapq では基本的に以下の 2 つがよく使われます。とりあえずこれだけ覚えておけば大体オッケーっぽいです。
heapq.heappush(hq, x) # リストhqに要素xを追加
heapq.heappop(hq) # リストhqから最小値を削除して返すheappush してすぐ heappop したいときなどは heappushpop を用いると若干計算量が削減されるようです。
heapq.heappushpop(hq, x) # リストhqに要素xを追加したのちhqの最小値を削除して返す最大値を取り出したいとき
heapq を使って最大値を取り出したいときには、要素全てを -1 倍して考えれば良いですね。
# a, b, c の最大を取り出したい
hq = []
heapq.heappush(hq, -a)
heapq.heappush(hq, -b)
heapq.heappush(hq, -c)
ans = -heapq.heappop(hq)
print(ans)出力するときに -1 倍してもとに戻すことを忘れずに。