【AtCoder参加記】PythonでABC249のA, B, C, Dを解く

自分がコンテストでやった考察と実装を書いてみます.

• A問題 『Jogging』
• B問題 『Perfect String』
• C問題 『Just K』
• D問題 『Index Trio』

A問題 『Jogging』

問題ページ: A - Jogging
Difficulty: 103

考察

難しくなかったですか?
for文で1秒毎に処理しました. 周期考えてもできるらしいです.
解説でもfor文でやってて, え? って思いました.

実装

#!/usr/bin/env pypy3

a, b, c, d, e, f, x = map(int, input().split())
res1 = 0
res2 = 0
for t in range(1, x + 1):
    t1 = t % (a + c)
    if 1 <= t1 <= a:
        res1 += b

    t2 = t % (d + f)
    if 1 <= t2 <= d:
        res2 += e

if res1 > res2:
    print("Takahashi")
elif res2 > res1:
    print("Aoki")
else:
    print("Draw")

B問題 『Perfect String』

問題ページ: B - Perfect String
Difficulty: 47

考察

大文字の文字全体の集合, 小文字の文字全体の集合を作っておくと, 条件は全て集合を使って $O(|S|)$ でできます.

実装

#!/usr/bin/env pypy3

from string import ascii_lowercase, ascii_uppercase

s = input()
s2 = set(s)
res = (len(s) == len(set(s)) and (s2 & set(ascii_uppercase))
       and (s2 & set(ascii_lowercase)))

if res:
    print("Yes")
else:
    print("No")

C問題 『Just K』

問題ページ: C - Just K
Difficulty: 528

考察

bit全探索の典型問題です. あんまり意識してなかったですけど, 久しぶりのbit全探索だったらしいですね.
s_1, s_2, ..., s_n から好きな個数を選ぶ場合の数は 2^n - 1 個です. これを全探索することを考えます.
それぞれに対して, 高々 $n$ 回和を計算すれば良いので, 計算量は $O(n 2^n)$ です.

実装

#!/usr/bin/env pypy3

n, k = map(int, input().split())
from collections import Counter

s = [None] * n
for i in range(n):
    s[i] = Counter(input())

from itertools import combinations

res = 0
for r in range(1, n + 1):
    for ts in combinations(s, r):
        cnt = sum(ts, start=Counter())
        c = 0
        for _, v in cnt.items():
            c += v == k
        res = max(c, res)

print(res)

D問題 『Index Trio』

問題ページ: D - Index Trio
Difficulty: 983

考察

$A_i$の最大値が $2 \cdot 10^5$ と小さめの値なのに注目して, この最大値についてループを回して計算できないかを考えます.
インデックス毎ではなく 値毎に計算するために cnt = Counter(A) を使って計算できないかを考えます.

問題は

A_i = A_j A_k

を満たす i, j, k の個数を求める問題だったので, 問題の答えば {xのA_i*A_jとして現れる回数} * cnt[x] のxに関する和です.
M = 2*10**5 と置く.
各z に対して {z = x*yのA_i*A_jとして現れる回数} を計算するために

for x in range(1, M + 1):
    for y in range(1, M+1):
       if x * y > 2 * M:
           break
       cnt2[x*y] = cnt[x] * cnt[y]

のようなコードを書く必要があります.
このコードの計算量を見てみましょう.
yは
x = 1 の時 M
x = 2 の時 M // 2
...
x=tの時 M // t
まで動くので, この処理の計算量は

M + M // 2 + M // 3 + \cdots + 1 = O(M \log M)

となるので, ACできます.

実装

#!/usr/bin/env pypy3

from collections import Counter

M = 2 * 10**5
n = int(input())
A = list(map(int, input().split()))
cnt = Counter(A)
cnt2 = Counter()

for x in range(1, M + 1):
    for y in range(1, M + 1):
        if x * y > M:
            break
        cnt2[x * y] += cnt[x] * cnt[y]

res = 0
for i in range(1, M + 1):
    res += cnt[i] * cnt2[i]

print(res)