Help us understand the problem. What is going on with this article?

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

はじめに

今回は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))

おわりに

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

Why do not you register as a user and use Qiita more conveniently?
  1. We will deliver articles that match you
    By following users and tags, you can catch up information on technical fields that you are interested in as a whole
  2. you can read useful information later efficiently
    By "stocking" the articles you like, you can search right away
Comments
Sign up for free and join this conversation.
If you already have a Qiita account
Why do not you register as a user and use Qiita more conveniently?
You need to log in to use this function. Qiita can be used more conveniently after logging in.
You seem to be reading articles frequently this month. Qiita can be used more conveniently after logging in.
  1. We will deliver articles that match you
    By following users and tags, you can catch up information on technical fields that you are interested in as a whole
  2. you can read useful information later efficiently
    By "stocking" the articles you like, you can search right away