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

2021.08.03

2024.03.24

競技プログラミング

AtCoderPython優先度付きキュー

本ページはAmazonアフィリエイトのリンクを含みます。

はじめに

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

優先度付きキューとは

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

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

heapqの使い方

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

1import heapq

ヒープの作成

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

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

要素の追加

heappushで要素を追加します。

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

最小の要素の取得

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

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

追加して取得

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

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

取得して追加

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

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

AtCoderのサンプル問題

ABC141 D問題

D - Powerful Discount Tickets

AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly programming contests online.

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

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

1import heapq
2
3n, m = map(int, input().split())
4A = list(map(lambda x: -int(x), input().split()))
5
6heapq.heapify(A)
7
8for _ in range(m):
9    min_price = heapq.heappop(A)
10    heapq.heappush(A, (-1)*(-min_price//2))
11
12print(-sum(A))

ABC173 D問題

D - Chat in a Circle

AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly programming contests online.

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

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

1import heapq
2
3n = int(input())
4A = list(map(int, input().split()))
5
6A.sort(reverse=True)
7
8ans = 0
9que = []
10heapq.heappush(que, -A[0])
11for i in range(1, n):
12    friendly = heapq.heappop(que) * (-1)
13    ans += friendly
14    heapq.heappush(que, -A[i])
15    heapq.heappush(que, -A[i])
16
17print(ans)

ABC212 D問題

D - Querying Multiset

AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly programming contests online.

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

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

1import heapq
2
3q = int(input())
4query = []
5for _ in range(q):
6    query.append(list(map(int, input().split())))
7
8bag = []
9heapq.heapify(bag)
10
11sum = 0
12for i in range(q):
13    if query[i][0] == 1:
14        heapq.heappush(bag, query[i][1]-sum)
15    elif query[i][0] == 2:
16        sum += query[i][1]
17    else:
18        ball_v = heapq.heappop(bag)
19        print(ball_v + sum)