【AtCoder】Pythonで優先度付きキュー

はじめに
優先度付きキューとは
heapqの使い方
AtCoderのサンプル問題
参考

はじめに

優先度付きキューについて、どんなものなのか、Pythonではどう使うか、実際のAtCoderの問題を解いてみたいと思います。

優先度付きキューとは

優先度付きキューとは、以下の操作が行えるデータ構造となります。

要素を追加する
最小の要素を取得・削除する

heapqの使い方

Pythonではheapqで優先度付きキューが利用できます。

import heapq

ヒープの作成

ヒープの作成は、初期化されたリスト[]かheapifyでリストを変換します。

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

要素の追加

heappushで要素を追加します。

que = []
heapq.heappush(que, 1)

最小の要素の取得

heappopで最小の要素を取得します。

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

追加して取得

heappushしてheappopするならheappushpopの方が効率が良いです。

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

取得して追加

heappopしてheappushするならheappushpopの方が効率が良いです。

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

AtCoderのサンプル問題

ABC141 D問題

D - Powerful Discount Tickets

AtCoderisaprogrammingcontestsiteforanyonefrombeginnerstoexperts.Weholdweeklyprogrammingcontestsonline.

N個の品物をM枚の割引券を使って買うとき、最小でいくら必要なのか計算する問題です。

割引券を1枚ずつ利用して、最も高い品物を割引していくことを繰り返します。最も高い品物を取り出すのに優先度付きキューを活用します。

import heapq

n, m = map(int, input().split())
A = list(map(lambda x: -int(x), input().split()))

heapq.heapify(A)

for _ in range(m):
    min_price = heapq.heappop(A)
    heapq.heappush(A, (-1)*(-min_price//2))

print(-sum(A))

ABC173 D問題

D - Chat in a Circle

AtCoderisaprogrammingcontestsiteforanyonefrombeginnerstoexperts.Weholdweeklyprogrammingcontestsonline.

N人ごとにフレンドリーさが決められており、環状に順番に並んでいく時に両隣の人のうちフレンドリーさが小さい方の心地よさを得ます。その心地よさを最大化する問題です。

並んだ人のフレンドリーさを優先度付きキューにプッシュしていきます。

import heapq

n = int(input())
A = list(map(int, input().split()))

A.sort(reverse=True)

ans = 0
que = []
heapq.heappush(que, -A[0])
for i in range(1, n):
    friendly = heapq.heappop(que) * (-1)
    ans += friendly
    heapq.heappush(que, -A[i])
    heapq.heappush(que, -A[i])

print(ans)

ABC212 D問題

D - Querying Multiset

AtCoderisaprogrammingcontestsiteforanyonefrombeginnerstoexperts.Weholdweeklyprogrammingcontestsonline.

袋に3種類の操作をしながら、ボールを取り出した時の数を出力していく問題です。

袋に入れたボールを優先度付きキューで管理します。

import heapq

q = int(input())
query = []
for _ in range(q):
    query.append(list(map(int, input().split())))

bag = []
heapq.heapify(bag)

sum = 0
for i in range(q):
    if query[i][0] == 1:
        heapq.heappush(bag, query[i][1]-sum)
    elif query[i][0] == 2:
        sum += query[i][1]
    else:
        ball_v = heapq.heappop(bag)
        print(ball_v + sum)

参考

<br/>