はじめに

今回はPythonの優先度付きキューについてまとめます。
AtCoderのPython3.4.3で動作確認済みです。

優先度付きキューについて

優先度付きキュー（Priority queue）はデータ型の一つで、具体的には

最小値（最大値）を $O(logN)$で取り出す
要素を $O(logN)$ で挿入する

ことが出来ます。通常のリストだとそれぞれ $O(N)$ ですので高速です。
「リストの要素の挿入」と「最小値（最大値）を取り出す」ことを繰り返すような時に使います。

Pythonでの使い方

Pythonでは優先度付きキューは heapq として標準ライブラリに用意されています。使いたいときはimportしましょう。

各メソッドについて

頻繁に使うメソッドは3つです。

heapq.heapify(リスト)でリストを優先度付きキューに変換。
heapq.heappop(優先度付きキュー)で優先度付きキューから最小値を取り出す。
heapq.heappush(優先度付きキュー)で優先度付きキューに要素を挿入。

import heapq  # heapqライブラリのimport

a = [1, 6, 8, 0, -1]
heapq.heapify(a)  # リストを優先度付きキューへ
print(a)

print(heapq.heappop(a))  # 最小値の取り出し
print(a)

heapq.heappush(a, -2)  # 要素の挿入
print(a)

出力

[-1, 0, 8, 1, 6]  # 優先度付きキューのa
-1                # aの最小値
[0, 1, 8, 6]      # 最小値を取り出した後のa
[-2, 0, 1, 8, 6]  # -2を挿入後のa

最大値の取り出し

heapqでは最小値しか取り出すことが出来ません。では最大値の時はどうするかというと、各要素に-1をかけた上で最小値を取り出していきます。
下のコードではmap関数で各要素を-1倍していますが、実際に問題を解く際には入力時に-1した方が当然高速です。

import heapq
a = [1, 6, 8, 0, -1]
a = list(map(lambda x: x*(-1), a))  # 各要素を-1倍
print(a)

heapq.heapify(a)
print(heapq.heappop(a)*(-1))  # 最大値の取り出し
print(a)

出力

[-8, -6, -1, 0, 1]
8
[-6, 0, -1, 1]

例題

AtCoder Beginners Contest 141 D問題 Powerful Discount Tickets
この問題は 最大値を1/2することをM回繰り返す だけでよいのですが、通常のリストを使うと、最大値の選定に $O(N)$ 、それをM回のため計算量が $O(NM)$ で間に合いません。
そこでheapqの出番です。heapqでは最大値の選定に $O(logN)$ 、それをM回のため計算量が $O(MlogN)$ となり間に合います。以下が実装です。

ABC141D.py

import heapq

n, m = map(int, input().split())
a = list(map(lambda x: int(x)*(-1), input().split()))
heapq.heapify(a)  # aを優先度付きキューに

for _ in range(m):
    tmp_min = heapq.heappop(a)
    heapq.heappush(a, (-1)*(-tmp_min//2))  # 計算誤差の関係で-1を2回かけてます
print(-sum(a))

おわりに

読んでいただきありがとうございました。指摘等ありましたら是非コメントお願いします。

【Python】優先度付きキューの使い方【heapq】