ABC324をPythonで解いてみたよ。(A～E問題)

AtCoder Beginners Contest 324 (ABC324) をPythonで解きました。
見やすいコードを書く練習も兼ねてます。

A - Same

問題ページ

考察

セット(集合)は、重複した値を持ちません。
なので、種類数を数えるときにつかえます。
この問題でも、配列 $A$ をセットとして受け取ってしまえば、簡単に判定ができます。

コード

N = int(input())
A = set(map(int, input().split()))
if len(A) == 1:
    print("Yes")
else:
    print("No")

B - 3-smooth Numbers

問題ページ

考察

入力値 $N$ が $2^x3^y$ の形をしているかどうかを判定するのもいいですが、
$2^x3^y$ の形で表せる数を先にたくさん列挙してしまって、その中に入力値 $N$ があるかどうかを判定すると、実装がラクでいいと思います。

制約は $N \leqq 10^{18}$ です。 $10^{18} < 2^{60}$ なので、 $0\leqq x,y \leqq 60$ の範囲で $2^x3^y$ の値をすべて列挙してセットに格納します。その中に $N$ があるかどうかを判定することで ACできます。

コード

se = set()
for i in range(61):
    for j in range(61):
        se.add((2 ** i) * (3 ** j))

N = int(input())
if N in se:
    print("Yes")
else:
    print("No")

C - Error Correction

問題ページ

考察

文字列 $T$ と $T´$ が同じかどうかは、次の条件で判断すると問題文にありました。

$T´$ は $T$ から一部が変更されてしまっている可能性があり、具体的には、下記の $4$ つのうちのちょうど $1$ つが成り立つことがわかっています。

• $T´$ は $T$ と等しい。
• $T´$ は、 $T$ のいずれか $1$ つの位置(先頭と末尾も含む)に英小文字を $1$ つ挿入して得られる文字列である。
• $T´$ は、 $T$ からある $1$ 文字を削除して得られる文字列である。
• $T´$ は、 $T$ のある $1$ 文字を別の英小文字に変更して得られる文字列である。

条件が $4$ つもあって、これを1つ1つ見ていくのはしんどいです。(頑張ればそれでもACできます)

上の $4$ つの条件のうち、 $1$ つ目と $4$ つ目は $T$ と $T´$ の文字数がおなじです。しかも条件もかなり似ていて、 $2$ つの文字列が全く一緒か、 $1$ 箇所だけ異なる文字になっているかです。 $1$ 文字目から $N$ 文字目まで順に見て、異なる文字の数が $1$ つ以下ならオッケー、そうでないならダメと判定すればよさそうです。

条件の$2$ つ目と $3$ つ目も実は似ていて、「長い方の文字列からある $1$ 文字を削除して短い方の文字列になる」かどうかで判定すればいいです。

コード

# 2つの文字列の長さがおなじとき
def judge1(A, B):
    diff_cnt = 0
    for a, b in zip(A, B):
        if a != b:
            diff_cnt += 1
    return diff_cnt <= 1


# 2つの文字列の長さが1つちがいのとき
def judge2(A, B):
    # 短い方をA, 長い方をBで固定にする。
    if len(A) > len(B):
        A, B = B, A
    j = 0
    for i in range(len(A)):
        if A[i] == B[j]:
            j += 1
            continue
        j += 1
        if i + 1 != j:
            return False
        if A[i] == B[j]:
            j += 1
        else:
            return False
    return True


N, T = input().split()
N = int(N)
ans_lst = []

for i in range(1, N + 1):
    S = input()
    if len(S) == len(T):
        if judge1(S, T): ans_lst.append(i)
    elif abs(len(S) - len(T)) == 1:
        if judge2(S, T): ans_lst.append(i)

print(len(ans_lst))
print(*ans_lst)

D - Square Permutation

問題ページ

考察

$0^2, 1^2, 2^2, \cdots$ と平方数を順番に見ていき、その平方数を並び替えて入力値の $S$ をつくれるかどうかを1つ1つ見ていきます。
桁数が足りない分は、先頭に $0$ をつけて桁数を揃えます。 s.zfill(N) で文字列 $s$ を $N$ 桁にするように先頭に $0$ をつけた文字列をゲットできます。
全体の個数が等しいかどうかはcollections.Counterをつかうと実装がラクです。

コード

from collections import Counter

N = int(input())
S = input()

max_num = 10 ** N
C = Counter(S)
sq = 0
ans = 0

while (num := sq * sq) < max_num:
    str_num = str(num).zfill(N)
    if Counter(str_num) == C:
        ans += 1
    sq += 1

print(ans)

E - Joint Two Strings

問題ページ

考察

たとえば、 $S_1=$"abba", $S_2=$"bcb" $T=$"bac" のとき、$S_1+S_2=$"abbabcb" となり、(連続してなくてもいい)部分文字列として "bac" があってokです。こうなるような$S_1, S_2$ のペアを求める問題です。

文字列 $S_i$ について、文字列 $T$ の先頭何文字が部分文字列になるか・文字列 $T$ の末尾何文字が部分文字列になるか、の$2$ つが分かればいいわけです。
先頭何文字が部分文字列になるかは、文字列 $S_i$ の先頭から貪欲に調べます。
末尾何文字が部分文字列になるかは、文字列 $S_i$ と文字列 $T$ をどちらも反転させて、「($S_i$を反転させた文字列)について、($T$を反転させた文字列) の先頭何文字が部分文字列になるか」を考えればいいです。
そして、(先頭 $x$ 文字が $T$ の部分文字列になる文字列 $S_1$ ) と (末尾 $y$ 文字が $T$ の部分文字列になる文字列 $S_2$) をくっつけたとき、 $x+y\geqq len(T)$ だったらokです。

各 $x$ について、 $len(T)-x$ 以上の $y$ がいくつあるかを調べればいいです。これは2分探索で求められます。

コード

from bisect import bisect_left

# 文字列Bの先頭何文字が文字列Aの部分文字列になっているか
def f(A, B):
    j = 0
    for i, a in enumerate(A):
        if a == B[j]:
            j += 1
            if j == len(B):
                break
    return j


N, T = input().split()
N = int(N)
S = [input() for _ in range(N)]

rev_T = T[::-1]
left = [0] * N
right = [0] * N
for i in range(N):
    left[i] = f(S[i], T)
    right[i] = f(S[i][::-1], rev_T)

right.sort()
ans = 0
for le in left:
    ans += N - bisect_left(right, len(T) - le)
print(ans)