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

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

A - Weak Beats

問題ページ

考察

文字列 $S$ の $2$ 番目、 $4$ 番目、 $6$ 番目、 $\cdots$ がすべて $0$ ならば答えはYesです。
逆に言えば、どれか1つでも $0$ ではなく $1$ になっていれば答えはNoです。この考え方で解いてみます。
Pythonでは $S$ の $2$ 番目 は S[1] 、 $4$ 番目なら S[3] と書くのに注意しましょう。

下のコードでは、for文で $i=0,1, \cdots ,15$ とループさせて、その中で $i$ が奇数のときに $S_i=1$ であればすぐにNoを出力して、exit()でコードを終了させています。

コード

S = input()
for i in range(16):
    if i % 2 == 1 and S[i] == '1':
        print("No")
        exit()
print("Yes")

別解

rangeの第3引数をつかった解法

$l$ から $r-1$ までの数を $s$ とばしで見るとき、次のように書きます。

"""例) 2～19の数を4つとばしで見る"""
for i in range(2, 20, 4):
    print(i)  # 2, 6, 10, 14, 18 が出力される 

これをつかうと、A問題は下のように書けます。

S = input()
for i in range(1, 16, 2):
    if S[i] != '0':
        print("No")
        exit()
print("Yes")

B - Round-Robin Tournament

問題ページ

考察

• 勝った回数の多い順
• 勝った回数が同じなら、プレイヤー番号の小さい順

のルールで並べる問題です。

$$G_x : x 回だけ勝ったプレイヤー番号のリスト$$
となるリスト $G$ を用意します。
プレイヤー $i$ の勝った回数、つまり $S_i$ の中に含まれるoの数は、S[i].count('o') などで数えられます。

あとは、$G_{N-1}$ から $G_0$ まで順に見て、各リストの中にあるプレイヤー番号を順番に出力していけばいいです。

下のコードでは、プレイヤー番号の小さいものから順に格納していっているので、 $G_0$ ～ $G_{N-1}$ すべてで小さいもの順に並んでいて、sort()メソッドはつかっていません。

コード

N = int(input())

G = [[] for _ in range(N)]
for p in range(1, N + 1):
    S = input()
    win_cnt = S.count('o')
    G[win_cnt].append(p)

for i in range(N - 1, -1, -1):
    for p in G[i]:
        # 小さいプレイヤー番号から順に格納していたので、
        # ソートしなくてもG[i]はすでに小さい順に並んでる。
        print(p, end=' ')

別解

(勝利回数, プレイヤー番号) のタプルでソートする解法

Pythonのソートは、基本的には小さいもの順に並べてくれます。
なので、勝利回数の多いもの順で並べるときは、勝利回数にマイナスの符号をつけることで小さいもの順に並べてokな形にします。

N = int(input())
S = [input() for _ in range(N)]

p = []
for i, s in enumerate(S):
    win_cnt = s.count('o')
    p.append((-win_cnt, i + 1))
p.sort()
for win_cnt, ans in p:
    print(ans, end =' ')

C - World Tour Finals

問題ページ

考察

まだ解いていない問題がいくつかある中で、配点の高い問題から順番に解いていきます。その方が解く問題数が少なくなるからです。
なので、どの順番で問題を解けばいいのかが分かるリストをあらかじめ作っておきます。

また、現状の各プレイヤーのポイントもあらかじめ計算しておきます。

あとは、各プレイヤーについて、先ほど作った解く順番のリストの通りに問題番号を見ていって、もしその問題を解いていなければその問題の配点だけポイントをプラスしていきます。
持っている得点が1位になった時点で、そのプレイヤーが問題を解いてポイントを増やす操作をやめます。

コード

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

B = [(-a, i) for i, a in enumerate(A)]
B.sort()  # # 点数の高い順

# 現状のポイントを計算する
now_point = [0] * N
for i, s in enumerate(S):
    now_point += i + 1
    for a, flag in zip(A, s):
        if flag == 'o':
            now_point[i] += a

max_point = max(now_point)
ans_lst = [0] * N
for i in range(N):
    point = now_point[i]
    if point == max_point:
        continue
    # 点数の高い順に、解いてない問題を解いていく
    for a, j in B:
        if S[i][j] == 'x':
            point += a
            ans_lst[i] += 1
        if point > max_point:
            break
for ans in ans_lst:
    print(ans)

D - Merge Slimes

問題ページ

考察

例えば、サイズ $12$ のスライムが $1$ つだけあったとします。
このスライムは現状だと合成できないですが、他にサイズ $6$ のスライムが $2$ つあったり、サイズ $3$ のスライムが $4$ つあったりすると、それを合成することでサイズ $12$ のスライムは $2$ つになって、サイズ $24$ のスライムを $1$ つつくれます。
サイズ $12$ だからまだ探索できましたが、問題の制約でサイズは $10^9$ 以下で、それくらいのサイズだといちいち探索していられません。

なので、「小さなスライムから順に見る」という方針が思いつきます。
これなら、サイズ $12$ のスライムが $1$ つあったときに、「サイズ $6$ のスライムが $2$ つあるんじゃないか...」と思わなくて済んでうれしいです。

「次々と新しい数が入ってくる中で、常に小さい順に数を取り出したい」わけです。これは、優先度付きキュー(ヒープキュー)で叶えられます。
ヒープキューにスライムのサイズを入れて、スライムの個数は辞書で管理します(キーがサイズ、バリューがそのサイズのスライムの個数)。
あとは、ヒープキューに要素が入っている間、「ヒープキューからサイズを取り出して、そのサイズのスライムを合成できるだけ合成してしまって、1つあまったら変数 $ans$ に $1$ を足す。」の操作を続けます。(下のコードも参考にしてください)

コード

from collections import defaultdict
from heapq import heappop, heappush

N = int(input())
dic = defaultdict(int)
que = []
for _ in range(N):
    s, c = map(int, input().split())
    dic[s] = c
    heappush(que, s)

ans = 0
while que:
    s = heappop(que)
    c = dic[s]
    nc, mo = c // 2, c % 2
    if nc:
        ns = s * 2
        if ns not in dic:
            heappush(que, ns)
        dic[ns] += nc
    ans += mo

print(ans)

E - Playlist

問題ページ

考察

動的計画法DPで解きます。
そもそも確率や期待値などをDPで解いたことがないよ！という人は、「期待値DP」でググってみて、EDPCやABCの過去問を解いてみてください。

• $dp[t]$ : 時刻 $t$ に音楽をかけ始める確率
• $ans$ : 時刻 $X+0.5$ に曲 $1$ が流れている確率

この2つの変数を、時刻 $0,1,\cdots,X$ で順に変形させていくだけです。
コード内にコメントを付けているので、詳しくは下のコードを見てください。

コード

MOD = 998244353

N, X = map(int, input().split())
T = list(map(int, input().split()))

inv = pow(N, MOD - 2, MOD)  # 1/n

# dp[t]: 時刻tに、音楽をかけ始める確率
dp = [0] * (X + 1)
dp[0] = 1

ans = 0  # 時刻X+0.5に、曲1が流れている確率

for t in range(X + 1):
    p = dp[t] * inv % MOD  # 時刻tに曲1をかけ始める確率(曲2,3,...も同じ)
    for i in range(N):
        # 曲iを流す場合を考える (これは確率pで起きる)
        nt = t + T[i]  # 曲iが流れ終わる時刻
        if i == 0 and X < nt:
            ans = (ans + p) % MOD
        if nt <= X:
            dp[nt] = (dp[nt] + p) % MOD
print(ans)

F - Push and Carry

問題ページ

考察

考えないといけないパターンは $4$ つだけです。

パターン1は、一般的なパターンです。
このとき、荷物とゴールの位置を考えると、プレイヤーは $(6,10)$ からとりあえず $(1,2)$ か $(2,1)$ まで移動します。
たとえば $(1,2)$ に移動したとします。
プレイヤーは、荷物を $(2,2)$ から $(10, 2)$ まで運びます。このときプレイヤーの位置は $(9,2)$ です。
次にプレイヤーは $(9,2) → (9,1) →(10,1)$ と動きます。 $\cdots (*)$
今度は荷物を上方向に押し上げます。 荷物を $(10,2) → (10,5)$ まで運びます。これでゴールです。

パターン2は、荷物とゴールのy座標がおなじになったパターンです。(x座標がおなじになったパターンも計算は一緒です)
このとき、プレイヤーは $(6,10)$ から $(1,5)$ まで移動します。
あとは荷物を $(2,5)$ から $(10,5)$ まで右方向にまっすぐ押すだけです。

パターン3は全部のy座標がおなじで、荷物が右端か左端にあるパターンです。(全部のx座標がおなじのときも計算は一緒です)
プレイヤーは $(6,5)→(6,4)→(1,4)→(1,5)$ と動きます。 $\cdots (*)$
あとは荷物を $(2,5) → (10,5)$ と動かして終わりです。

パターン4は、荷物が真ん中にあるパターンです。
プレイヤーは $(2,5)→(9,5)$ と荷物を動かしながら移動します。

どのパターンも、「荷物のあるマスに隣接したマスまで移動する」→「荷物をゴールまで運ぶ」の流れです。
パターン1とパターン3は、荷物がジャマで迂回するシーン$\cdots (*)$があって、実装のときは気をつけましょう。

コード

xa, ya, xb, yb, xc, yc = map(int, input().split())

if xa == xb == xc:
    xa, ya, xb, yb, xc, yc = ya, xa, yb, xb, yc, xc

pos = []
if xb < xc:
    pos.append((xb - 1, yb))
elif xb > xc:
    pos.append((xb + 1, yb))
if yb > yc:
    pos.append((xb, yb + 1))
elif yb < yc:
    pos.append((xb, yb - 1))

dist2 = abs(xb - xc) + abs(yb - yc)
f1 = xb != xc and yb != yc  # パターン1
f2 = ya == yb == yc and not (xa < xb < xc or xc < xb < xa)  # パターン3
if f1 or f2:
    dist2 += 2  # 迂回する分

ans = 1 << 80
for x, y in pos:
    dist1 = abs(x - xa) + abs(y - ya)
    ans = min(ans, dist1 + dist2)
print(ans)