ABC306をPythonで解いてみたよ。(A～E問題)

AtCoder Beginners Contest 306 (ABC306) をPythonで解きました。
見やすいコードを書く練習も兼ねてます。

A - Echo

問題ページ

考察

文字を1つずつ見ていって、その文字を2回出力します。
Pythonのprint関数はデフォルトだと改行が入ってしまうので、おまじないで第2引数にend=''とつけてあげます。

print関数のendについて

先に結論から書くと、endのところには出力した後の末尾を書きます。 たとえばこんな感じ。

# 末尾にxをつけて出力するならこう。
for i in range(4):
    print(i, end='x')  # 0x1x2x3x

# よくあるのは下の2つ。
# 改行なし、全部ぴったりくっつけて出力
for i in range(4):
    print(i, end='')  # 0123

# 改行なし、半角スペースをあけて出力
for j in range(4):
    print(i, end=' ')  # 0 1 2 3

コード

N = int(input())
S = input()
for s in S:
    print(s + s, end='')

B - Base2

問題ページ

考察1

2進数で考えます。
たとえばサンプル1は答えが$13$になっていて、2進数表記すると、$13=(1101)_2$ です。
$A_0$が1の位、$A_1$が10の位、...となっていて、普段見ている数字と左右が逆になっているのに注意します。
リストAの左右を反転させて、それを2進数だと思って10進数表記に直してあげればACです。

コード1

A = list(input().split())
A = A[::-1]   # 左右の反転でよくつかいます！
N = ''.join(A)  # リストを1つの文字列にする。
print(int(N, 2))  # 2進数の文字列を10進数の数に変換する。

考察2

こっちの方が簡単で、問題文の通りに式を立てます。
(Pythonだとオーバーフローが起こらなくて、これはPython有利(?)な問題だと思います！)

コード2

A = list(map(int,input().split()))

ans = 0
for i in range(64):
    ans += A[i] * (2 ** i) 
print(ans)

C - Centers

問題ページ

考察

それぞれの数字について、何番目に現れたのかを記録します。
具体的には、$A[i]=num$のときリスト$G[num]$に$i$を格納します。
同じ数字は3回出てくるらしいので、どの$i$についても$G[i]$の中には3つしか数字が入っていません。
問題文で出てくる$f(i)$は、$G[i][1]$ のことになります。
$f(i)$の小さいものから順に$i$を出力すればいいので、リストPを用意して、その中に$(f(i), i)$のタプルを格納していって、最後にPを昇順にソートすればいいです。

コード

N = int(input())
A = list(map(int, input().split()))

G = [[] for _ in range(N)]

for i, a in enumerate(A):
    G[a - 1].append(i)

P = []
for i in range(N):
    P.append((G[i][1], i + 1))
P.sort()
for a, b in P:
    print(b, end=' ')

D - Poisonous Full-Course

問題ページ

考察

動的計画法DPで解きます。次のように2次元リストdpを決めます。
$dp[i][flag]$: $i$個目までの品物が出てきて、高橋くんの状態がflag(0なら元気、1なら腹痛)のときの、食べた料理のおいしさの総和
初期状態は $dp[0][0]=0$ です。

先に出された料理を食べないときを考えます。次のように遷移します。
$dp[i+1][0]=dp[i][0]$
$dp[i+1][1]=dp[i][1]$

次に出された料理を食べる時を考えます(おいしさを$v$とします)。
料理に毒が盛られていないとき、
$dp[i+1][0]=max(dp[i+1][0], dp[i][0]+v, dp[i][1]+v)$
料理に毒が盛られているとき、腹痛だと食べられないことに注意します。
$dp[i+1][1]=max(dp[i+1][1], dp[i][0]+v)$

答えは$max(dp[N-1][0], dp[N-1][1])$です。

コード

INF = 1 << 80

N = int(input())
P = [list(map(int, input().split())) for _ in range(N)]

# dp[i][flag]: max_score
# flag=0: 元気、 flag=1: やばい
dp = [[-INF] * 2 for _ in range(N + 1)]
dp[0][0] = 0

for i in range(N):
    dp[i + 1] = dp[i][:]
    x, v = P[i]
    if x == 0:
        dp[i + 1][0] = max(max(dp[i]) + v, dp[i + 1][0])
    else:
        dp[i + 1][1] = max(dp[i][0] + v, dp[i + 1][1])

print(max(dp[-1]))

E - Best Performances

問題ページ

考察

tatyamさん作のSortedMultiSetをつかいます(いつもありがとうございます！)。
大きい方から上位k個を入れておくSortedMultiSetと、それ以外のSortedMultiSetをつくります。
また、1つ目のSortedMultiSetに入ってる数の総和を表す変数totも用意しておきます。

あとは実際にシミュレートするのみです。
シミュレートするのに上の3つが必要だということが分かりますか？というところが1つ目の壁で、これを時間内にコードにできますか？というのが2つ目の壁な気がします。

どうシミュレートするかを書くとかなり長くなるので、詳しくはコードを見てください。

コード

'''
tatyamさん作の、SortedMultisetです。
使わせていただき、ありがとうございます！
https://github.com/tatyam-prime/SortedSet/blob/main/SortedMultiset.py

SortedSetの多重集合版。同じ要素を複数入れることができる。
SortedSetから、s.add(x), s.discard(x) が変更され、s.count(x)が追加されている。

・使い方(個人的まとめ)
s=SortedMultiSet()
s.a: SortedMultiSetの中身を返す。
len(s), x in s, x not in s: リストと同じ要領で使える。
s.add(x): xを追加する。sにxが含まれているかどうかは関係ない。
s.discard(x): xを1つだけ削除してTrueを返す。もし存在しないなら、Falseを返す。
s.count(x): sに含まれるxの個数を返す。
s.lt(x): xより小さい最大の要素を返す。もし存在しないなら、Noneを返す。
s.le(x): x　以下の　最大の要素を返す。もし存在しないなら、Noneを返す。
s.gt(x): xより大きい最小の要素を返す。もし存在しないなら、Noneを返す。
s.ge(x): x　以上の　最小の要素を返す。もし存在しないなら、Noneを返す。
s.index(x): xより小さい要素の数を返す。
s.index_right(x): x以下の要素の数を返す。

・使い方URL
https://github.com/tatyam-prime/SortedSet
'''
# (SortedMultiSetが長いから省略。上のURLからコピペしてね。)

N, K, Q = map(int, input().split())

A = [0] * N
use_set = SortedMultiset()
for _ in range(K):
    use_set.add(0)
gar_set = SortedMultiset()  # garbage
for _ in range(N - K):
    gar_set.add(0)
tot = 0

for _ in range(Q):
    x, y = map(int, input().split())
    prev = A[x - 1]
    A[x - 1] = y
    c = gar_set.discard(prev)

    # gar_setからprevを取り出しているとき
    if c:
        if use_set.lt(y) is not None:
            el = use_set.ge(-1)
            use_set.discard(el)
            gar_set.add(el)
            tot -= el
            use_set.add(y)
            tot += y
        else:
            gar_set.add(y)

    # gar_set内にprevがなかったとき
    else:
        use_set.discard(prev)
        tot -= prev
        if use_set.le(y) is None:
            gar_set.add(y)
            el = gar_set.le(1000000007)
            gar_set.discard(el)
            use_set.add(el)
            tot += el
        else:
            use_set.add(y)
            tot += y
    print(tot)