リストの中から一番小さい値を取り出したい。しかも、値を足しながら何度も取り出したい。
そんな場面で、ループのたびにsortを呼んでいませんか。動くには動きますが、データが増えると少しずつ重くなっていきます。
Pythonには、こうした用途のための道具が標準で入っています。それがheapqです。
名前だけ聞くと難しそうですが、使う関数はほんの数個です。今回は、仕組みのイメージから実務での使いどころまで、動くコードと一緒に見ていきます。
heapqは、最小値を取り出すことに特化した道具¶
heapqは、ヒープと呼ばれるデータの並べ方を扱うモジュールです。ヒープは、一番小さい値がいつも先頭にいることだけを保証した、ゆるい並べ方だと考えてください。
ふつうの並べ替えは、すべての要素を小さい順にきっちり揃えます。ヒープはそこまでせず、先頭が最小であることだけを守ります。
全部を揃えないからこそ、出し入れが速い。これがヒープのいちばん大事な性質です。
ちなみに、heapqは専用のクラスを作りません。ふつうのPythonのリストを、ヒープのルールに従って並べ直して使います。
ヒープの中身は、なぜバラバラに見えるのか¶
実際にリストをヒープにしてみると、少し不思議な並びになります。まずは次のコードを動かしてみましょう。
import heapq
scores = [72, 95, 58, 88, 64, 91, 70]
heapq.heapify(scores)
print(scores) # [58, 64, 70, 88, 95, 91, 72]
print(scores[0]) # 58 → 先頭が必ず最小値
先頭の58は最小値ですが、その後ろは小さい順になっていません。これで正しい状態です。
ヒープは、リストを木の形に見立てています。位置kの要素は、位置2k+1と2k+2の要素以下である、というルールだけが守られていればよいのです。
ですから、ヒープにしたリストをそのままprintして、並んでいないと慌てる必要はありません。見るべきは先頭の値だけです。
基本の関数は、heappushとheappopの2つ¶
heapqで覚えたい関数は多くありません。よく使うものを表にまとめると次のとおりです。
| 関数 | やること |
|---|---|
| heapify(x) | リストxをその場でヒープに並べ直す |
| heappush(heap, item) | ヒープのルールを保ったまま値を追加する |
| heappop(heap) | 最小値を取り出して返す |
| heap[0] | 最小値を取り出さずに見るだけ |
| heappushpop(heap, item) | 追加してから最小値を取り出す |
| heapreplace(heap, item) | 最小値を取り出してから追加する |
| nsmallest(n, data) / nlargest(n, data) | 小さい順または大きい順に上位n件を返す |
| merge(*iterables) | 並べ替え済みのデータ同士を順番を保って合体する |
heappushpopとheapreplaceは、押し込みと取り出しを1回でまとめた関数です。別々に呼ぶより効率がよいので、両方を続けて行う場面では思い出してください。
空のリストから始めるときは、heapifyは要りません。空のリストは、それだけで正しいヒープだからです。
空のヒープからpopするとエラーになる¶
ひとつだけ注意があります。要素がないヒープでheappopを呼ぶと、IndexErrorが発生します。
取り出す前に、while heapのようにリストが空でないかを確かめる書き方が定番です。空のリストは偽として扱われるので、このままで条件式になります。
なぜ毎回sortするより速いのか¶
ここで、少しだけ計算量の話をさせてください。データの件数をnとしたとき、処理の手間がどう増えるかを表す考え方です。
sortは、全体を並べ替えるのにおよそn log nの手間がかかります。値を1つ足すたびにsortし直すと、この手間を毎回払うことになります。
一方で、heappushとheappopは1回あたりlog n程度で済みます。公式ドキュメントでも、既存のリストをheapifyでヒープにする処理は線形時間、つまりnに比例する手間だと説明されています。
たとえば100万件のデータなら、log nは20ほどにしかなりません。何度も出し入れするほど、差ははっきり開いていきます。
計算量という考え方そのものに自信がない方は、こちらで基礎から解説しています。【関連記事】「実行時間が終わらない…」を卒業する！あなたのコードを100倍速くする計算量の考え方
上位N件を取るならnlargestとnsmallest¶
ヒープを自分で管理しなくても、heapqの便利さを味わえる関数があります。nlargestとnsmallestです。
売上の上位3件や、レスポンスが遅かったAPIのワースト10を出したい。そんなときに1行で書けます。
import heapq
scores = [72, 95, 58, 88, 64, 91, 70]
print(heapq.nlargest(3, scores)) # [95, 91, 88]
print(heapq.nsmallest(2, scores)) # [58, 64]
users = [
{"name": "佐藤", "age": 31},
{"name": "鈴木", "age": 24},
{"name": "高橋", "age": 45},
]
oldest = heapq.nlargest(1, users, key=lambda u: u["age"])
print(oldest) # [{'name': '高橋', 'age': 45}]
sortedと同じように、key引数で比べる基準を指定できます。辞書やオブジェクトのリストでも、そのまま使えるのがうれしいところです。
key引数の考え方は、sortedの記事で詳しく扱いました。【関連記事】Pythonのsortedとsortの違いとは？key引数で自在に並べ替える方法を初心者向けに解説
sortedやmaxとの使い分け¶
では、いつでもnlargestを使えばよいのでしょうか。実はそうとも限りません。
公式ドキュメントでは、nが小さいときにこの2つの関数が力を発揮すると説明されています。取り出したいのが1件だけならminやmaxのほうが速く、nが全体に近いならsortedで並べて切り出すほうが効率的だとも書かれています。
| 欲しいもの | おすすめ |
|---|---|
| 最大値や最小値を1つだけ | max() / min() |
| 全体に比べて少ない上位N件 | heapq.nlargest() / nsmallest() |
| ほぼ全件を順番どおりに | sorted()してスライス |
| 値を足しながら何度も最小値を取り出す | heappush() / heappop() |
迷ったときは、何度も出し入れするかどうかで考えてみてください。一度きりならsortedで十分なことも多いです。
優先度付きキューを作ってみる¶
heapqのいちばん代表的な使い道が、優先度付きキューです。入れた順番ではなく、優先度の高いものから取り出す行列のことです。
たとえば、障害対応は最優先、掃除は後回し、というタスクの管理を考えてみます。優先度を小さい数字ほど高いと決めれば、heapqがそのまま使えます。
import heapq
from itertools import count
tasks = []
counter = count() # 同じ優先度のときの順番を記録する
def add_task(priority, name):
heapq.heappush(tasks, (priority, next(counter), name))
add_task(2, "メール送信")
add_task(1, "障害対応")
add_task(2, "レポート作成")
add_task(3, "掃除")
while tasks:
priority, _, name = heapq.heappop(tasks)
print(priority, name)
# 1 障害対応
# 2 メール送信
# 2 レポート作成
# 3 掃除
ポイントは、タプルの2番目に通し番号を入れていることです。タプルは先頭から順に比べられるので、優先度が同じときは先に入れたものが先に出てきます。
通し番号を入れないと、どうなるのか¶
通し番号を省いて、優先度と名前だけのタプルにしても一見動きます。ですが、同じ優先度のものが並ぶと、名前の文字列同士が比べられてしまいます。
名前ではなく辞書や自作クラスを入れていると、比べ方が決まっていないためTypeErrorになります。公式ドキュメントも、この問題を避けるために追加の番号を挟む方法を紹介しています。
私は10年ほど開発に関わってきましたが、この落とし穴を本番で踏んだことがあります。テストでは優先度がたまたま重ならず、本番のデータで初めて同じ優先度の辞書同士が比べられてエラーになりました。
それ以来、ヒープに何かを入れるときは、最初から通し番号を挟む書き方に決めています。優先度が重なるケースを前提に書く。小さな習慣ですが、夜中の障害対応を1つ減らしてくれます。
タプルの代わりにクラスで管理したい場合は、dataclassのorder=Trueを使う方法もあります。【関連記事】Pythonのdataclassとは？クラスの定義がぐっと短くなる書き方を初心者向けに解説
大きい順に取り出したいとき¶
heapqの関数は、どれも最小値を取り出す前提で作られています。では、最大値から順に取り出したいときはどうすればよいでしょうか。
昔からよく使われてきたのは、値にマイナスを付けて入れる方法です。マイナスを付けると大小が逆転するので、最小値として取り出したものが元の最大値になります。
import heapq
heap = []
for x in [5, 1, 8]:
heapq.heappush(heap, -x)
print(-heapq.heappop(heap)) # 8
取り出したあとに、もう一度マイナスを付けて元に戻すのを忘れないでください。数値でしか使えないのも、この方法の弱点です。
Python 3.14からは、最大ヒープ用の関数が使える¶
この不便さを解消するため、Python 3.14では最大ヒープ用の関数が公開されました。heapify_max、heappush_max、heappop_max、heappushpop_max、heapreplace_maxの5つです。
使い方は、最小ヒープ用の関数名の末尾に_maxを付けるだけです。3.14以降の環境なら、次のように素直に書けます。
import heapq # Python 3.14以降
heap = []
for x in [5, 1, 8]:
heapq.heappush_max(heap, x)
print(heapq.heappop_max(heap)) # 8
ただし、3.13以前ではこれらの関数は使えません。チームやサーバーのPythonのバージョンを確かめてから使いましょう。
heapqが活躍する場面¶
ここまでの内容を、実際の使いどころと結びつけておきます。私が現場やコードレビューで見かけた例をまとめると、次のようになります。
| 場面 | heapqの使い方 |
|---|---|
| ジョブやタスクを優先度順に処理する | 優先度付きキュー |
| ログから遅かったリクエストの上位を出す | nlargest |
| 日付順に並んだ複数のCSVを1本にまとめる | merge |
| 最短経路を求めるダイクストラ法 | 距離が最小の地点を順に取り出す |
| 大量データから上位N件だけを保持し続ける | サイズNのヒープで入れ替え |
最短経路の問題は、競技プログラミングでも定番です。AtCoderなどに挑戦していると、heapqを使う問題に必ずと言ってよいほど出会います。【関連記事】競技プログラミング（AtCoder）は実務に役立つ？Pythonで挑戦するメリットとデメリット
merge関数は、すでに並んでいるデータ同士を1本にするときに便利です。全部を読み込んでsortし直すより、メモリを節約できる場面があります。
つまずきやすいところ¶
最後に、初心者がひっかかりやすい点をまとめておきます。
ひとつめは、ヒープにしたリストを、並べ替え済みのリストだと思い込むことです。先頭以外は順番どおりではないので、scores[1]が2番目に小さいとは限りません。
ふたつめは、ヒープにしたリストに、appendで直接値を足してしまうことです。ヒープのルールが崩れると、heappopが正しい最小値を返さなくなります。追加は必ずheappushで行ってください。
みっつめは、スレッドをまたいで共有することです。複数のスレッドから同時に出し入れするなら、ロックを自分で用意するか、標準ライブラリのqueue.PriorityQueueを検討しましょう。
まとめと、次の一歩¶
heapqは、最小値を何度も取り出すという仕事に特化した道具です。全部を並べ替えないからこそ、データが増えても軽く動きます。
まずはheappushとheappop、そして上位N件を取るnlargestとnsmallestを覚えておけば十分です。優先度付きキューでは、通し番号を挟むことを忘れないでください。
次の一歩として、手元のコードでループの中にsortが書かれていないか探してみてください。毎回sortしている場所は、heapqに置き換えられる候補です。
道具を知っているかどうかで、同じ処理の書き方は大きく変わります。今日の内容が、あなたのコードを少し軽くするきっかけになればうれしいです。
ここまでお読みいただきありがとうございました。
