テクめも

プログラミング関連のちょっとしたTipsなどを書いています。

Pythonで優先度付きキューを使う

Pythonでの優先度付きキューはheapqを使うとできます。

import heapq

最小ヒープ

はじめに、リストをキューに優先度付きキューに変換します。

que = [3, 1, 5, 2]
heapq.heapify(que)

先頭に最小値がきます。(先頭以外はソートされているわけではありません。)

print(que)
# [1, 2, 5, 3]

要素を追加します。(計算量: O(logN)

heapq.heappush(que, 4)
print(que)
# [1, 2, 5, 3, 4]

heapq.heappush(que, -1)
print(que)
# [-1, 1, 2, 5, 3, 4]

要素を取り出します。(計算量: O(1)

v = heapq.heappop(que)
print(v)
# -1

最大ヒープ

heapqでは最小ヒープですが、最大ヒープを利用したい場合は、デフォルトではないので-1をかける必要があります。

que = [4, 5, 9, 7]
que = [x * -1 for x in que]
heapq.heapify(que)

print(que)
# [-9, -7, -4, -5]
heapq.heappush(que, 2 * -1)
print(que)
# [-9, -7, -4, -5, -2]

v = heapq.heappop(que) * -1
print(v)
# 9