読者です 読者をやめる 読者になる 読者になる

Pythonでヒープ構造を使う(heapqモジュール)

常に最初の要素が最小値(あるいは最大値)となるようなリストが必要な場合、ヒープ構造を用いることで最小値(最大値)の取り出しをO(1)、要素の追加をO(log n)の時間計算量で行うことができる。

Pythonでヒープを扱う場合、heapqモジュールが使える。

import heapq

a = [6, 3, 2, 4, 5]
heapq.heapify(a)    # 破壊的

print a[0]    # => 2
heapq.heappop(a)    # 最小値をヒープから取り出す
print a[0]    # => 3
heapq.heappush(a, 1)    # ヒープに要素を追加する
print a[0]    # => 1

heapqは最初の要素が最小値となるようなヒープしか扱えない。最初の要素が最大値となるようなヒープ(max-heap)を使う場合は、下記のようにタプルで包んで扱うか、__lt__()の戻り値が本来の結果と逆になるようなクラスを作って扱う必要がある。

a = [6, 3, 2, 4, 5]
a = [(-x, x) for x in a]
heapq.heapify(a)    # 破壊的

print a[0][1]    # => 6
heapq.heappop(a)
print a[0][1]    # => 5
heapq.heappush(a, (-8, 8))
print a[0][1]    # => 8

なんかイマイチ。