【ABC173 ABCD】優先度付きキュー(ヒープ)を使えば簡単【Python】

本記事について

C問題は「総探索」，
D問題は「優先度付きキュー」，
を使って回答した

A問題「Payment」

1000円札だけで買い物したら，いくらのお釣りが出るかという問題．
難しく考えず，問題文そのままって感じ．

n=int(input())
for i in range(1000,10001,1000):
　　if i>=n:
　　　　print(i-n)
　　　　break

B問題「Judge Status Summary」

ただの数え上げ問題．
連想配列(辞書型)を使えば，難なく書ける．

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

result={"AC":0,"WA":0,"TLE":0,"RE":0}

for i in s:  #数え上げ
　　result[i]+=1

for k,v in result.items():  #出力
　　print("{} x {}".format(k,v))

C問題「H and V」

問題内容については以下を参照
https://atcoder.jp/contests/abc173/tasks/abc173_c

この問題で考えることは以下の2点である

赤く塗る行と列
黒いマス(“#”)の数え上げ

1に関しては，行Hと列Wが「1<=H,W<=6」という制約下なので，総探索で解いても間に合いそうだと直感的にわかる．

2に関しても総探索で判定できそう．

まずは，2に関するソースコードを書いてみる．

# 黒("#")の数え上げ
# H:行，W:列
cnt=0
for row in range(H):
　　for col in range(W):
　　　　if c[row][col]=="#":
　　　　　　cnt+=1

次に，上記のソースコードに1の条件を追加する．
考え方としては，赤く塗る場合を1，赤く塗らない場合を0として，起こり得る全ての組み合わせ(0と1)を列挙する．(以下のソースコードと例)

rows = itertools.product(range(2),repeat=H)

例：H=2の場合，rows = [ [0,0], [0,1], [1,0], [1,1] ]が列挙
[0,0] = どの行も赤く塗らない
[0,1] = 1行目を赤く塗る
[1,0] = 2行目を赤く塗る
[1,1] = 1,2行目を赤く塗る

また，黒マスは赤く塗られることで黒マスではなくなるので気を付ける．
これらの条件を組み合わせて書いたのが，以下のソースコード．

from itertools import product


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

H,W,K=input2()
c=[list(input()) for _ in range(H)]

ans=0 

for row_line in product(range(2),repeat=H): #[行]0:そのまま，1:赤く塗る
　　for col_line in product(range(2), repeat=W): #[列]0:そのまま，1:赤く塗る
	# 黒("#")の数え上げ
	cnt=0
	for row in range(H):
	　　for col in range(W):
		if c[row][col]=="#" and row_line[row]==0 and col_line[col]==0:
		　　cnt+=1

	if cnt==K: #黒いマスがちょうどK個である
	　　ans+=1
print(ans)

D問題「Chat in a Circle」

問題内容については以下を参照
https://atcoder.jp/contests/abc173/tasks/abc173_d

考え方のステップとして以下の通り．

入力配列Aを降順にソート
“心地良さ”が最大になる位置に値を挿入していく

図で表すと，以下のような画像になる．

赤文字が次の”心地良さ“の値を表している．
この値が最も大きい場所に次の人が割り込めば良い．
つまり，この赤文字部分を配列の中に加えていき，その中での最大値を取り出していく．
今回は最大値を効率良く取り出すために，「優先度付きキュー」を利用して実装した．

結果，以下のソースコードで通った．

import heapq

#入力:[n1,n2,...nk](int:整数配列)
def input_array():
	return list(map(int,input().split()))

n=int(input())
A=input_array()
sortA=sorted(A,reverse=True) #降順にソート

ans=0
q=[]
heapq.heappush(q,-sortA[0]) #初期値

for i in range(1,n):
	tmp=heapq.heappop(q) 
	tmp*=-1
	ans+=tmp
	heapq.heappush(q,-sortA[i]) #挿入した値の左右に次の人が割り込む場所が生まれる
	heapq.heappush(q,-sortA[i])
print(ans)