ABC471のPython解答（A～E）

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

A問題

調べる

A

a, b = map(int, input().split())

if 9 in [a + b, a - b, a * b, a / b]:
    print("Nine")
else:
    print("Nein")

B問題

upperにそろえてからカウントして最大個数を求める

B

n = int(input())

d = dict()
for _ in range(n):
    s = input().upper()
    if s not in d:
        d[s] = 0
    d[s] += 1

print(max(d.values()))

C問題

プラスとマイナスでソートして現在地から近いほうを順番に見ていく

C

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

plus, minus = list(), list()
for a_i in a:
    if a_i >= 0:
        plus.append(a_i)
    else:
        minus.append(a_i)

plus.sort()
minus.sort(reverse=True)
ans = 0
min_ind = max_ind = 0
now = 0
while min_ind < len(minus) and max_ind < len(plus):
    if abs(minus[min_ind] - now) <= abs(plus[max_ind] - now):
        ans += abs(minus[min_ind] - now)
        now = minus[min_ind]
        min_ind += 1
    else:
        ans += abs(plus[max_ind] - now)
        now = plus[max_ind]
        max_ind += 1

if min_ind < len(minus):
    ans += abs(minus[-1] - now)
elif max_ind < len(plus):
    ans += abs(plus[-1] - now)

print(ans)

D問題

$W_q-T_q$でheapqに入れて最大値を取得する

D

from heapq import heappop, heappush

def max_pop(q):
    res = heappop(q)
    return -res

def max_push(q, target):
    heappush(q, -target)


q, v = map(int, input().split())

lst = list()

for _ in range(q):
    com = list(map(int, input().split()))
    if com[0] == 1:
        t, w = com[1:]
        max_push(lst, w - t)
    else:
        t = com[1]
        if len(lst) <= 0:
            print(-1)
        else:
            x = max_pop(lst)
            print(min(v, x + t))

E問題

ずばり答えを式変形すると以下になる

\sum_{i=1}^n( a_i^2\times {}_{n-1}C_{k-1} + a_i(sum(A) - a_i) \times {}_{n-2}C_{k-2} )

E

def cmb(n, r, mod):
    from functools import reduce

    numerator = reduce(lambda x, y: x * y % mod, [n - r + k + 1 for k in range(r)])
    denominator = reduce(lambda x, y: x * y % mod, [k + 1 for k in range(r)])
    return numerator * pow(denominator, mod - 2, mod) % mod


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

sum_a = sum(a) % mod
ans = 0
if k == 1:
    for a_i in a:
        ans += pow(a_i, 2, mod)
        ans %= mod
elif k == 2:
    p = cmb(n - 1, k - 1, mod)

    for a_i in a:
        ans += pow(a_i, 2, mod) * p % mod + a_i * (sum_a - a_i) % mod
        ans %= mod
else:
    p1 = cmb(n - 1, k - 1, mod)
    p2 = cmb(n - 2, k - 2, mod)

    for a_i in a:
        ans += pow(a_i, 2, mod) * p1 % mod + a_i * (sum_a - a_i) % mod * p2 % mod
        ans %= mod

print(ans)