AtCoder ABC222 個人的メモ

A - Four Digits

文字列Nの長さが4以上になるまで、Nの先頭に0をつけ続ける。

N = input()

while len(N) < 4:
    N = "0" + N

print(N)

B - Failing Grade

全ての学生の点数を見て、P点未満かどうかを判定する。

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

ans = 0
for a in A:
    if a < P:
        ans += 1
    
print(ans)

C - Swiss-System Tournament

問題文通りにシミュレーション

from operator import itemgetter

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

# (勝数,人のインデックス)が順位の低い順に並んでいる
ranking = [(0, i) for i in range(N * 2)]
for j in range(M):
    new_rank = []
    while ranking:
        a_win_cnt, a_i = ranking.pop()
        b_win_cnt, b_i = ranking.pop()
        a = A[a_i][j]
        b = A[b_i][j]

        # aがジャンケンに勝ったかどうかを判定
        if a == "G" and b == "C":
            is_a_win = 1
        elif a == "G" and b == "P":
            is_a_win = 0
        elif a == "C" and b == "G":
            is_a_win = 0
        elif a == "C" and b == "P":
            is_a_win = 1
        elif a == "P" and b == "G":
            is_a_win = 1
        elif a == "P" and b == "C":
            is_a_win = 0
        else:
            is_a_win = -1

        # aが勝った場合
        if is_a_win == 1:
            new_rank.append((a_win_cnt + 1, a_i))
            new_rank.append((b_win_cnt, b_i))
        # aが負けた場合
        elif is_a_win == 0:
            new_rank.append((a_win_cnt, a_i))
            new_rank.append((b_win_cnt + 1, b_i))
        # あいこの場合
        else:
            new_rank.append((a_win_cnt, a_i))
            new_rank.append((b_win_cnt, b_i))

    # 各人の番号でソート → 勝数でソート とすると問題の条件での順位順にソートできる
    new_rank.sort(key=itemgetter(1))
    new_rank.sort(reverse=True, key=itemgetter(0))
    ranking = list(new_rank)

for _, i in ranking:
    print(i + 1)

参考

pythonのソートの安定性の話
https://docs.python.org/ja/3/howto/sorting.html#sort-stability-and-complex-sorts

D - Between Two Arrays

dp

N = int(input())
A = list(map(int, input().split()))
B = list(map(int, input().split()))
MOD = 998244353
max_val = 3000

# dp[i][j]:=数列のi番目まででC_{i-1}=jとなる場合の数
dp = [[0] * (max_val + 1) for _ in range(N + 1)]
dp[0][0] = 1
for i in range(1, N + 1):
    a, b = A[i - 1], B[i - 1]
    cumsum = sum(dp[i - 1][:a]) % MOD
    for j in range(a, b + 1):
        cumsum += dp[i - 1][j]
        cumsum %= MOD
        dp[i][j] = cumsum

print(sum(dp[-1]) % MOD)

E - Red and Blue Tree

https://atcoder.jp/contests/abc222/editorial/2751

"""
頂点をAの順番で巡った時に各辺を通る回数を数える
これをR-B=Kとなるように赤と青に割り振る
    数字がN-1個あって、これを2つのグループに分けた時にそれぞれの総和の差がKになる組み合わせはいくつ？
    dpでつくれる数を数える → Kとなる組み合わせの個数を求める
その塗り方が解
"""
from collections import deque


def bfs(start: int, end: int):
    """
    頂点startから頂点endまでBFSする
    後に経路復元するためにhistoryに各頂点にどの頂点からやって来たかをメモ
    """
    queue = deque([start])
    while queue:
        now = queue.popleft()
        if now == end:
            return

        for n_node in edge[now]:
            if history[n_node] != -1:
                continue
            history[n_node] = now
            queue.append(n_node)
    return


N, M, K = map(int, input().split())
A = list(map(lambda n: int(n) - 1, input().split()))
edge = [[] for _ in range(N)]
edge_trans = {}  # 頂点x-y間の辺の辺番号を返す
for i in range(N - 1):
    x, y = map(lambda n: int(n) - 1, input().split())
    edge[x].append(y)
    edge[y].append(x)
    x, y = sorted([x, y])
    edge_trans[(x, y)] = i
MOD = 998244353

visit_cnt = [0] * N
for i in range(1, M):
    # historyは頂点iにどの頂点から来たかを記録
    history = [-1] * N
    bfs(A[i - 1], A[i])
    # 経路を復元してどの辺を通ったかをvisit_cntに記録
    now = A[i]
    while now != A[i - 1]:
        n_node = history[now]
        edge_i = edge_trans[tuple(sorted([now, n_node]))]
        visit_cnt[edge_i] += 1
        now = n_node

# dp[i]:=(B=iとなる場合の数)
dp = [0] * (10 ** 5 + 1)
dp[0] = 1
for val in visit_cnt[:N - 1]:
    for i in range(len(dp) - 1 - val, -1, -1):
        if dp[i]:
            dp[i + val] += dp[i]

ans = 0
sum_all = sum(visit_cnt)
for num, v in enumerate(dp):
    R = sum_all - num
    B = num
    if R - B == K:
        ans += v
        ans %= MOD

print(ans)