【ABC137 A~D問題】優先度付きキューを使って解いてみる(解答とコード)【Python3】

本記事について

AtCoder Beginner Contest 第137回のA~D問題を，Python3で解いてみた．
https://atcoder.jp/contests/abc137

特にD問題は「優先度付きキュー」というアルゴリズムを使えば解けるらしいので，それを使って解いてみた．

A問題「+-x」

説明はほとんど不要だと思うので，省略．

#入力:N,M(int:整数)
def input2():
	return map(int,input().split())

a,b=input2()
ans=a+b
if a-b>ans:
	ans=a-b
if a*b>ans:
	ans=a*b 
print(ans)

B問題「One Clue」

変数 st = (黒石の開始地点)
変数 fi = (黒石の終了地点)

st~fiの位置に置かれている黒石を配列 ans に格納して出力．

#入力:N,M(int:整数)
def input2():
	return map(int,input().split())

k,x=input2()

st=x-k+1
fi=x+k-1
ans=[]
for i in range(st,fi+1):
	ans.append(i)

print(*ans)

ちょこっとポイント(配列出力編)

ちなみに，配列 ans の要素を横並びに出力したい時，「*ans」と書けばいいらしい．
他の方法だと，「(” “).join(ans)」という方法もある．

C問題「Green Bin」

AtCoderだと茶色の平均的なレベル(?)の問題．

入力された文字列の中で，アナグラム(入れ替え文字列)が等しい文字列の組の合計を出力する．
今回は入力文字列をすべてソートして，辞書に格納し，アナライズの出現回数を記録することにした．

[8~9行目]：入力された文字列をソート(a~z)し，変数 ss とする
[10~14行目]：変数 ss が辞書に未登録なら登録し，登録済なら出現数を更新し，ansに加算する

n=int(input())
s=[]
for _ in range(n):
	s.append(str(input()))

counter={}
ans=0
for i in range(n):
	ss = sorted(s[i])
	ss=("").join(ss)
	if ss not in counter:
		counter[ss]=0
	else:
		counter[ss]+=1
		ans+=counter[ss]

print(ans)

D問題「Summer Vacation」

Atcoderだと水色レベルの問題．
問題内容は以下の通り．

B日後に報酬Aが得られるアルバイトが何個かある．
M日後までに得られる報酬が最大となるようなアルバイトを選択し，こなしていく．
結果，報酬はいくらになったか．

これは初日はこれ，2日目はこれという風に選択すると，めんどくさい．
なので，

M-1日目(M日後の1日前)は，Bが1のアルバイトからAが最大のものを選択．
M-2日目(M日後の2日前)は，Bが1もしくは2のアルバイトからAが最大のものを選択．
．
．
．
1日目(初日)は，残っているアルバイトからAが最大のものを選択

という感じで選んでいけば良い．

ここで重要なのは，選択可能なアルバイトから報酬が最大のものを選択，つまり最大値を効率よく取り出すアルゴリズムが必要ということである．

そこで使用するのが，「優先度付きキュー」である．

優先度付きキューとは

アルゴリズムの一つである「ヒープ(heap)」を少し応用した手法．
データ数がN個であるデータに対して，

データの挿入をO(logN)で行う
最大値(最小値))の取り出しをO(logN)で行う

ことができる．

以下，参考URL
https://ufcpp.net/study/algorithm/col_heap.html

Python3では，「heapq」というモジュールがあるので，それを利用すると素早く実装できる．

import heapq

#入力:N,M(int:整数)
def input2():
	return map(int,input().split())

n,m=input2()
AB=[[] for _ in range(m)]


for i in range(n):
	a,b=input2()
	if a-1 < m: #報酬受取日がmを超える仕事は除外する
		AB[a-1].append(-b) #heapqの仕様のため，-を追加

ans=0
heap=[]
for i in range(m):
	for b in AB[i]:
		heapq.heappush(heap,b)
	#最大値をheapから取り出し
	if len(heap)>0:
		MAX=heapq.heappop(heap)
		ans+= -MAX
print(ans)