Python heapqの使い方

heapqとは

Pythonの標準ライブラリの一つで、優先度付きキュー(priority queue)の実装です。

本記事では、heapqという表現で統一します。

heapqの特徴

最小値の取得が高速

heapqを用いた最小値の取得を計算量O(1)で行えます。これはとても高速です。

なぜなら、組み込み関数min()は計算量O(N)だからです。

これだけ聞くと、heapqを用いるのが常に最適に思えますが、そうではありません。

heapqにはいくつか前提条件があります。

heapqの前提条件

あるリストが与えられたとしましょう。

1
A = [3,5,1,2,0]

このリストの最小値を求めるために、まずはheapの構築を行う必要があります。

1
heapq.heapify(A)

Aの中身を確認しましょう。

1
2
>>> A
[0, 2, 1, 3, 5]

順番が入れ替わっていますね。

リストAはヒープと呼ばれるデータ構造(正確には、Pythonではある規則に従って並び替えられたリスト)となっています。

ただし、ソートされているわけではありません。

この時

1
heap[0]

で最小値を取得することができます。
Pythonのリストは、計算量O(1)で要素へアクセスできます。

そして、先ほどのヒープの構築heapq.heapifyには計算量O(N)を要します。

従って、ここまで合計すると計算量O(N)で最小値を取得できたことになります。

そのため、heapqでは最小値を取得するために計算量O(1)のコストで済むが、
その前提として、ヒープの構築(コストはO(N))が必要ということです。

つまり、ただ最小値を得たいだけであれば

組み込み関数min()を使っても、heapqを使っても計算量は変わらないですし、

(恐らく、オーバーヘッドがある分、heapqが遅い)

記述量も多くなるため、メリットはありません。

heapqを活用すべき場面

heapqが威力を発揮するのは小さい方から順に複数の値を取得したい場合です。

heapqでリストから小さい順にK個の値を取り出したい場合

heapq.heappopを使って、最小値を削除していく必要があります。

実装例は以下の通りです。

1
2
3
4
5
A = [3,5,1,2,0]
heapq.heapify(A)
K = 2
for k in range(K):
print (heapq.heappop(A))

のように

heappopの計算量はO(logN)であるため、計算量はheapqの構築も含めてO(N+KlogN)となります。

取り出す回数が小さければheapifyの計算量O(N)が支配的と理解しましょう。

ソートとの比較

比較手法として、ソートを使う実装が考えられます。

例えば、リストから小さい順に2つの値を取り出したい場合、以下のように実装すれば良いです。

1
2
A = [3,5,1,2,0]
sorted(A)[:2]

sorted()の計算量はO(NlogN)です。

heapqを使うべき場面 まとめ

ソートされていないリストから小さい順にK個の値を取り出したい場合の計算量は

  • sort: O(NlogN)
  • heapq: O(N+KlogN)

となるので、Kの数(取り出したい数)に応じて使い分けるべきです。

  • K=1の場合: min
  • KがNより十分小さいの場合: heapq
  • KとNがほぼ等しい場合: sort

ただし、後述するように、数を取り出す途中で挿入も発生する場合はheapqを使うべきです。

最大値を取得したい場合

これを実現するヒープを最大ヒープと言います。

heapqの実装では最小値を取得するような仕様となっています。(最小ヒープ)
ですので、heapify前のリストやheappushする際の変数に、-1をかけることで最大ヒープを実現できます。

最小ヒープ・最大ヒープを、一つのヒープで実現することはできないので注意してください。
(無理矢理やるのであれば、切り替えの都度リスト全体に-1をかけheapifyをコールすることになると思います。)

数を取り出す途中で挿入が発生する場合

hepaqは計算量O(logN)で挿入を行うこともできます。

1
heapq.heappush(A, 5)

ソート済みのリストへの挿入は、標準ライブラリbisectで実装されています。二分探索自体はO(logN)で実現できますが挿入の計算量がO(N)
です。そのためこのような場合もheapqを用いると良いでしょう。

最後に

Pythonでのstackqueueの使い方については以下の記事も参考にしてください。

Pythonでstackとqueueを使う

記事情報

  • 投稿日:2020年3月18日
  • 最終更新日:2020年5月16日