【AtCoder】Pythonで優先度付きキュー
はじめに
優先度付きキューについて、どんなものなのか、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)
参考
- heapq --- ヒープキューアルゴリズム — Python 3.9.4 ドキュメント
- 解説 - AtCoder Beginner Contest 141
- 解説 - AtCoder Beginner Contest 173
- 解説 - AtCoder Beginner Contest 212