ABC357をPythonで（A～E）

サントリープログラミングコンテスト2024（AtCoder Beginner Contest 357）の解答等のまとめ

A問題

先頭から順番に足していく

A

n, m = map(int, input().split())
a = list(map(int, input().split()))

s = 0
for i, a_i in enumerate(a):
    s += a_i
    if s > m:
        print(i)
        exit()

print(n)

B問題

問題文の通りにする

B

s = input()

ups = sum(s_i.isupper() for s_i in s)
if len(s) <= ups * 2:
    print(s.upper())
else:
    print(s.lower())

C問題

再帰関数で書く

C

def dfs(n):
    if n == 0:
        return ["#"]
    else:
        a = dfs(n - 1)
        b = ["." * 3 ** (n - 1) for _ in range(3 ** (n - 1))]
        res = []
        for a_i in a:
            # print(a_i)
            res.append(a_i + a_i + a_i)
        for a_i, b_i, a_j in zip(a, b, a):
            # print(a_i, b_i, a_j)
            res.append(a_i + b_i + a_i)
        for a_i in a:
            res.append(a_i + a_i + a_i)
        return res


ans = dfs(int(input()))
for a_i in ans:
    print(a_i)

D問題

初項$1$公比$10^m$（$m$は$n$の桁数）項数$n$の等比級数の和に$n$をかける

D

n = int(input())
mod = 998244353
m = len(str(n))

a = pow(10, m, mod)
print((pow(a, n, mod) - 1) * pow(a - 1, -1, mod) % mod * n % mod)

E問題

各頂点からの到達可能な頂点の数は

• ループしているところはそのループ内の頂点数
• それ以外は行けるループ内の頂点数＋そこまでの距離

E

import sys

sys.setrecursionlimit(10 ** 7)

n = int(input())
a = list(map(lambda x: int(x) - 1, input().split()))

loop_list = []
memo = [0] * n


def loop(now, turn):
    global memo, loop_list
    memo[now] = turn
    to = a[now]
    if now == to:
        loop_list.append({now})
        return False
    elif memo[to] == 0:
        flag = loop(to, turn)
        if flag:
            if now in loop_list[-1]:
                return False
            else:
                loop_list[-1].add(now)
                return True
        else:
            return False
    elif memo[to] == turn:
        loop_list.append({to, now})
        return True
    else:
        return False


cnt = 1
for i in range(n):
    if memo[i] == 0:
        loop(i, cnt)
        cnt += 1

ans = [0] * n
for l_i in loop_list:
    cnt = len(l_i)
    for i in list(l_i):
        ans[i] = cnt


def dfs(now):
    to = a[now]
    if ans[to] == 0:
        dfs(to)
    ans[now] = ans[to] + 1


for i in range(n):
    if ans[i] == 0:
        dfs(i)

print(sum(ans))