ABC444をPythonで

AtCoder Beginner Contest 444の解答等の速報的まとめ

A問題

文字列としてすべて同じか調べる

A

print("Yes" if len(set(list(input())))==1 else "No")

B問題

そのままシミュレーション

B

n, k = map(int, input().split())
ans = 0
for i in range(1, n + 1):
    s = list(str(i))
    if sum(map(int, s)) == k:
        ans += 1

print(ans)

C問題

元の本数$k$の条件は

• 長さの合計の約数
• $n/2 \leq k \le n$

よってそれを満たすパターンすべてでペアが作れるか調べたらいい
以下のコードはCodonで提出

C

from collections import Counter

def calc_divisor(x):
    res = list()
    i = 1
    while i * i <= x:
        if x % i == 0:
            res.append(i)
            if x // i != i:
                res.append(x // i)
        i += 1
    res.sort()
    return res

def check(target):
    for key in sorted(C.keys()):
        if key == target:
            continue
        else:
            pair = target - key
            if key == pair and C[key] % 2 > 0:
                return False
            elif pair not in C:
                return False
            elif C[key] != C[pair]:
                return False
    return True


n = int(input())
a = [int(x) for x in input().split()]

sum_a = sum(a)
ans = list()
div = calc_divisor(sum_a)
C = Counter(a)

for k in div:
    if n < k:
        break
    if k * 2 < n:
        continue
    target = sum_a // k
    if check(target):
        ans.append(target)

ans.sort()
for ans_i in ans:
    print(ans_i, end=" ")

D問題

各桁をimos法で計算

D

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

dp = [0] * (max(a) + 10)
for a_i in a:
    dp[a_i - 1] += 1

for i in range(len(dp) - 1, 0, -1):
    dp[i - 1] += dp[i]

ans = list()
i = 0
while i < len(dp):
    ans.append(str(dp[i] % 10))
    if dp[i] >= 10:
        if len(dp) >= i:
            dp.append(0)
        dp[i + 1] += dp[i] // 10
    i += 1

while len(ans) > 1 and ans[-1] == "0":
    ans.pop()

print("".join(ans[::-1]))

E問題

SortedListを使って尺取り法で調べる

E

from sortedcontainers import SortedList

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

s = SortedList([])
ans = 0
right = 0
for left in range(n):
    while right < n:
        ind = s.bisect_left(a[right])
        if ind - 1 >= 0 and abs(a[right] - s[ind - 1]) < d:
            break
        if ind < len(s) and abs(a[right] - s[ind]) < d:
            break
        s.add(a[right])
        right += 1
    ans += right - left

    s.discard(a[left])

print(ans)