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

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

A - 3.14

問題ページ

考察

問題文に円周率の小数点以下を書いてくれているので、それをコピペして文字列として受け取ります。
たとえば小数第2位までを聞かれたら"3.14"と答えます。文字列としてみると4文字です。
小数第 $N$ 位までを聞かれたら、答えは先頭 $N+2$ 文字になってます。

コード

S = "3.1415926535897932384626433832795028841971693993751058209749445923078164062862089986280348253421170679"
N = int(input())
print(S[:N + 2])

B - Roulette

問題ページ

考察

出力する人の番号のリストansと、その人たちの書けていた番号の個数cntの2つの変数をつかいます。
人1から順番に見ていきます。
もし $X$ に賭けていなかったら、その人は対象外です。 $X$ に賭けていたかどうかはsetをつかって判定できます。
次にcntを見ます。もし賭けていた番号の個数がcntと同じなら、ansにその人の番号を追加します。 もし賭けていた番号の個数がcntよりも少ないなら、ansを初期化してからこの人の番号を追加してあげて、あとcntもこの人の賭けていた番号の個数に更新しておけばokです。

コード

N = int(input())
P = []
for _ in range(N):
    c = int(input())
    A = set(map(int, input().split()))
    P.append(A)

X = int(input())

ans = []
cnt = 1000
for i in range(N):
    if X in P[i]:
        if cnt == len(P[i]):
            ans.append(i + 1)
        elif cnt > len(P[i]):
            ans = [i + 1]
            cnt = len(P[i])

print(len(ans))
print(*ans)

C - Rotate Colored Subsequence

問題ページ

考察

色1で塗ったものだけをピックアップしてシフトして、色2で塗ったものだけをピックアップしてシフトして、...としたいです。
塗る色ごとに、「この色で塗ってるインデックスはどこかなー」と探していると $O(MN)$ の計算量でTLEしちゃいます。
これは、先に

$D_i$ : 色 $i$ で塗られたインデックスのリスト

になっている2次元リスト $D$ をつくってあげると解消できる問題です。

下のコードでは、答えとなる文字列 $T$ を別で用意してあげて、入力で受け取った文字列 $S$ は変更しないようにしています。こうしなくても解けるかもしれないんですが、個人的にこれがラクでこうしています。

コード

N, M = map(int, input().split())
S = list(input())
C = list(map(int, input().split()))

D = [[] for _ in range(M)]
for i, c in enumerate(C):
    D[c - 1].append(i)

T = S.copy()
for gr in D:
    if len(gr) <= 1:
        continue
    for i in range(len(gr)):
        prev_idx = gr[i]
        nex_idx = gr[(i + 1) % len(gr)]
        prev = S[prev_idx]
        T[nex_idx] = prev
print(*T, sep='')

D - LOWER

問題ページ

考察

$t=2$ や $t=3$ のクエリ(全部を小文字にしたり大文字にしたり)をクエリが飛んでくるごとに素直に処理していると、計算量 $O(NQ)$ でTLEしてしまいます。
$t=1$ のクエリは文字を1つ変えるだけでそんなに時間がかからなさそうです。問題になっているのは $t=2,3$ のクエリっぽいです。
では、最悪のパターンでクエリが $t=2, t=3, t=2, t=3,...$ と交互にとんできたときにどうするかを考えてみます。
結論から書くと、$t=2$ もしくは $t=3$ のクエリのうち、最後に来たものだけを採用すればいいです。それ以外のクエリは全部無視してしまって問題ないです(問題のサンプル1とサンプル2で実験してみると分かりやすいと思います)。

クエリがくるごとに処理せずに、先にすべてのクエリを受け取ってしまいます。このクエリたちの中で最後に来た $t=2$ もしくは $t=3$ がいつなのかを先に記録しておいて、それ以外の $t=2$ もしくは $t=3$ のクエリは全部無視すればokです。(「クエリ先読み」と呼ばれる頻出テクニックです)

コード

N = int(input())
S = list(input())

last_idx = -1

Q = int(input())
query = []
for qi in range(Q):
    t, x, c = input().split()
    t, x = int(t), int(x) - 1
    query.append((t, x, c))
    if t == 2 or t == 3:
        last_idx = qi

for qi in range(Q):
    t, x, c = query[qi]
    if qi == last_idx:
        if t == 3:
            for i in range(N):
                S[i] = S[i].upper()
        elif t == 2:
            for i in range(N):
                S[i] = S[i].lower()
        continue
    if t == 1:
        S[x] = c

print(*S, sep='')

E - Roulettes

問題ページ

考察

期待値DPと呼ばれる類の問題です。
次のリストを用意します。

dp[i]: iポイントを持っていて、Mポイント以上にするために必要な金額の期待値

このリストを末尾から順番に完成させていきます。期待値DPは後ろから解いてみる、というのは典型です。(例： EDPC-J Sushi)

具体的には、たとえばサンプル1を解いていて、 $dp[5]$ ～ $dp[M]$ は分かっていて $dp[4]$ を求めたいとします。
どのルーレットを使うのが最適かは分からないので、全部試してみて、一番期待値のよかったルーレットを採用することにします。
まずコスト $100$ 円で目が $(5,9)$ のルーレットを使ってみます。$5$ の目が出たら $4+5=9$ ポイントになり、 $9$ の目が出たら $4+9=13$ ポイントになります。なので、このルーレットをつかうと、$dp[4]=100+(dp[9]+dp[13])/2$ 円になります。
これと同じことを他のルーレットでもすべてやってみて、支払う金額の期待値が最も低かったものを採用すればokです。

また、 $0$ が含まれているルーレットに関しては、たとえばサンプル2だと $(0,0,0,0,0,100)$ のルーレットがありますが、$100$ の目が出る確率は $1/6$ で、 $100$ の目を出すために必要なルーレットを回す回数の期待値は $6$ 回です。
一般的に、確率 $p$ で起こる事象があったとして、その事象を起こすのに必要な回数の期待値は $1/p$ になります。($1$ %しかあたらないガチャを当てるのに、平均で $100$ 回ガチャを回さないといけないね、的なことです)
$0$ の目が含まれるすべてのルーレットについて、$0$ の目をなくしてその分だけ代わりにコストを増やしておきます。たとえば先ほどの $(0,0,0,0,0,100)$ でコスト $10$ 円のルーレットなら、 $(100)$ のルーレットに変換して、 $0$ でない目が出る確率は $1/6$ だったので、コストを $6$ 倍して $60$ 円にします。

コード

N, M = map(int, input().split())
R = []
C = []
for _ in range(N):
    c, p, *s = map(int, input().split())
    zero_cnt = s.count(0)
    s.sort()
    if zero_cnt:
        c *= len(s) / (len(s) - zero_cnt)
        s = s[zero_cnt:]
    C.append(c)
    R.append(s)

# dp[i]: iポイントを持っていて、Mポイント以上にするために必要なコインの期待値
dp = [1 << 80] * (M + 1)
dp[M] = 0

for i in range(M - 1, -1, -1):
    for j in range(N):
        now_val = 0
        cost = C[j]
        rou = R[j]
        p = 1 / len(rou)
        for el in rou:
            nex = min(M, el + i)
            now_val += (dp[nex] + cost) * p
        dp[i] = min(dp[i], now_val)
print(dp[0])