Python の heapq について調べたのでメモ【AtCoder】【Python3】

heapq

先日の ABC137 の D 問題を解く際に、 Python の heapq について軽く調べたのでメモを残しておきます。

そもそも heapq ってなんやねーん

先にざっくり概要を説明すると、 heapq というのは python のモジュールで、プライオリティキュー(優先度付きキュー)というデータ構造を実現するときに利用されます。

heapq は次の操作を何度も繰り返すようなときに強さを発揮します。

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 heapq

heapq を活用するためにまず 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 倍してもとに戻すことを忘れずに。

この記事を書いた人

サトゥー

東大学際情報学府M1。情報科学と教養の海に溺れています。面白いことをやるのがすきです。