【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

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

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

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)

参考

Support

\ この記事が役に立ったと思ったら、サポートお願いします! /

buy me a coffee
Share

Profile

author

Masa

都内のIT企業で働くエンジニア
自分が学んだことをブログでわかりやすく発信していきながらスキルアップを目指していきます!

buy me a coffee