ABC347をPythonで解いてみたよ。(A～E問題+おまけ)

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

A - Divisible

問題

正整数 $N, K$ と、数列 $A=(A_1, A_2, \cdots ,A_N)$ が与えられます。
$A$ の要素のうち $K$ の倍数であるものをすべて取り出し、それらを $K$ で割った値を出力してください。

考察

答えとなるリスト $B$ を用意します。
「 $K$ の倍数である」という条件は、「 $K$ で割ったときの余りが $0$ になる」ことと同じです。
for文をつかって数列 $A$ の各要素を見ていきます。
$A$ の要素 $a$ について、 $K$ で割ったときの余りが $0$ であれば、 $a/K$ をリスト $B$ に挿入していきます。

リストの出力は print(*B) と書くとラクです。

コード

N, K = map(int, input().split())
A = list(map(int, input().split()))

B = []
for a in A:
    if a % K == 0:
        B.append(a // K)

print(*B)

別解

リスト内包表記をつかった解法

リスト内包表記をつかうと、少し簡潔に書けます。

N, K = map(int, input().split())
A = list(map(int, input().split()))

B = [a // K for a in A if a % K == 0]
print(*B)

B - Substring

問題

英小文字からなる文字列 $S$ が与えられます。
$S$ の連続した部分文字列は何種類ありますか？

考察

制約を見ると、文字列 $S$ の長さは最大でも $100$ です。
連続した部分文字列の取り方のパターン数は $_{N+1}{C}_2$ あり、最大で $5050$ 通りです。
そこまでパターン数が多くないので、部分文字列をすべてチェックする全探索で解くことができます。

種類数を求めるとき、重複した文字列をまとめて $1$ つとしてカウントしたいので、setをつかいます。
Pythonでは、$S$ の $i$ 文字目から $j-1$ 文字目までを取り出した部分文字列を S[i:j] と書きます。 けっこう便利な書き方なので、知らなかった方はこの機会に覚えてしまいましょう。

コード

S = input()

N = len(S)
se = set()
for i in range(N):
    for j in range(i + 1, N + 1):
        se.add(S[i:j])

print(len(se))

C - Ideal Holidays

問題

$1$ 週間の長さが $A+B$ 日で、$1$ 日目から $A$ 日目は休日、 $A+1$ 日目から $A+B$ 日目が平日であるとします。
今日が $1$ 週間の何日目かは分かりません。あなたは予定が $N$ 個あり、 $i=1, 2, \cdots ,N$ について、$i$ 個目の予定は今日から $D_i$ 日後です。
これら $N$ 個の予定がすべて休日であることはありえますか？

考察1 (数列Dを変換する)

現実世界の $1$ 週間は $7$ 日間です。
たとえば、今日から $2$ 日後、 $9$ 日後、 $16$ 日後、 $\cdots$ はすべて同じ曜日のはずで、どれかが休日で別のどれかが平日になることはありません。

それは $1$ 週間が $A+B$ 日間からなるこの問題設定でも同じで、「$d$ 日後が休日かどうか」と「$d+A+B$ 日後が休日かどうか」は等しくなるはずです。

したがって $i=1, 2, \cdots ,N$ について、$D_i$ を $A+B$ で割ったあまりに置き換えます。
またこの操作によって重複した要素が出てくることがあります。何度も同じ日付を確認する必要はないので、重複した要素を入れないようにsetで管理してしまいます。

N, A, B = map(int, input().split())
D = list(map(int, input().split()))

P = set()
for d in D:
    P.append(d % (A + B))

また、次の考察のために、このsetをソート済みのリストに変換しておきます。

P = list(sorted(P))

考察2 (問題を言い換える)

たとえば、休日が $6$ 日間、平日が $4$ 日間の下のカレンダーについて考えます。
(オレンジが休日、緑が平日です。)

休日の最後の日(=曜日 $3$) から休日の最初の日(=曜日 $8$) の間に、平日が $4$ 日間あります。
一般的に、休日の最後の曜日が $x_1$ で休日の最初の曜日が $x_2$ ($x_1<x_2$) のとき、$x_2-x_1-1=B$ が成り立つので、$x_2-x_1 \gt B$ が成り立ちます。

これを考察1で用意した予定のある曜日リスト $P$ と合わせて考えます。
$P_0$ と $P_1$ について $P_1-P_0 \gt B$ だったとき、曜日$P_0$ の翌日から曜日 $P_1$ の前日までの間に $B$ 日間の平日があったとすると、$N$ 日の予定はすべて休日になります。
これは $(P_0, P_1)$ の組だけでなく、 $(P_1, P_2), (P_2, P_3), (P_3, P_4), \cdots$ の組についても同様です。

これを実装すると下のコードになります。

($(P_{|P|-1}, P_0)$ の組については、$P_{|P|-1} \gt P_0$ なので実装に気をつけましょう。下のコードでは $P_{|P|-1} - P_0$ を $A+B$ で割ったあまりで考えることでこの問題を解決しています)

コード

N, A, B = map(int, input().split())
D = list(map(int, input().split()))

P = {d % (A + B) for d in D}
lst = list(sorted(P))

if len(lst) == 1:
    print("Yes")
    exit()

for i in range(len(lst)):
    d1 = lst[i]
    d2 = lst[(i + 1) % len(lst)]
    if (d2 - d1) % (A + B) > B:
        print("Yes")
        exit()
print("No")

D - Popcount and XOR

問題

$3$ つの非負整数 $a, b, C$ が与えられます。以下の条件をすべて満たす整数組 $(X, Y)$ はありますか？
あるなら $X, Y$ を出力してください。なければ $-1$ と出力してください。

• $\text{popcount}(X) = a$
• $\text{popcount}(Y) = b$
• $0 \leq X, Y \lt 2^{60}$
• $X \oplus Y = C$

考察1 (全体の方針)

$2$ 進数表記で $C$ の下から $k$ 桁目が $0$ のとき、$X,Y$ の下から $k$ 桁目はどちらも $1$ かどちらも $0$ です。
$C$ の下から $k$ 桁目が $1$ のとき、$X, Y$ の下から $k$ 桁目は片方が $1$ でもう片方は $0$ です。

$C$ の下から $k$ 桁目が $0$ のときについて、$X,Y$ の下から $k$ 桁目を $0$ にするか $1$ にするかを選べて、これで $X, Y$ の $\text{popcount}$ が足りていないときに調整できます。
それに比べて $C$ の下から $k$ 桁目が $1$ のとき、$X, Y$ の下から $k$ 桁目について片方を $1$ もう片方を $0$ にしないといけないので、$\text{popocount}$ の総量は変わりません。

なので、先に $C$ の $2$ 進数表記で $1$ になっている桁について計算してから、あとで $\text{popcount}$ を調整する形で $C$ の $2$ 進数表記で $0$ になっている桁について考えることにします。

考察2 (実装する)

先に $C$ の $2$ 進数表記で $1$ になっている桁について考えます。
その桁について、 $X, Y$ の片方を $1$ に、もう片方を $0$ にするのでした。
$a-\text{popcount}(X)$ と $b- \text{popcount}(Y)$ を比較して、大きい方に $1$ を、小さい方に $0$ を割り振っていきます。

コードにするとこんな感じ。

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

x, y = 0, 0
px, py = 0, 0  # px: popcount(x)

for k in range(60):
    if (C >> k) & 1 == 1:
        if a - px >= b - py:
            x |= 1 << k
            px += 1
        else:
            y |= 1 << k
            py += 1

次に、$C$ の $2$ 進数表記で $0$ になっている桁について考えます。
この桁について、 $X, Y$ どちらも $0$ にするかどちらも $1$ にするかを選べるのでした。

$\text{popcount}(X)<a$ かつ $\text{popcount}(Y)<b$ のときはどちらも $1$ にして、そうでないときはどちらも $0$ にすることにします。

最後に $\text{popcount}(X) = a, \text{popcount}(Y) = b$ となっていればokです。

以上をまとめると、下のコードになります。

コード

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

x, y = 0, 0
px, py = 0, 0  # px: popcount(x)

for k in range(60):
    if (C >> k) & 1 == 1:
        if a - px >= b - py:
            x |= 1 << k
            px += 1
        else:
            y |= 1 << k
            py += 1

for k in range(60):
    if (C >> k) & 1 == 0:
        if px < a and py < b:
            x |= 1 << k
            y |= 1 << k
            px += 1
            py += 1
if (a, b) == (px, py):
    print(x, y)
else:
    print(-1)

E - Set Add Query

問題

数列 $A=(A_1, A_2, \cdots ,A_N)$ があります。最初、 $A_1=A_2=\cdots =A_N=0$ です。
また集合 $S$ があります。最初、集合 $S$ は空です。
$Q$ 個のクエリが与えられます。このクエリをすべて処理した後の $A$ を出力してください。
$i=1, 2, \cdots Q$ について、 $i$ 個目のクエリは次のものです。

• $x_i$ が与えられる。もし $S$ 内に $x_i$ があれば $S$ から $x_i$ を取り除き、そうでなければ $S$ に $x_i$ を追加する。 その後、$j=1, 2, \cdots ,N$ について、 $j$ が $S$ 内にあるとき、 $A_j$ に $|S|$ を加算する ($|S|$ は $S$ の要素数)

考察

いくつかシミュレーションを行うと、次の事実が見えてきます。

• $i$ 回目のクエリ $x_i$ について、 $1$ 回目のクエリからここまでで $x_i$ が登場したのが今回で奇数回目のとき、集合 $S$ に $x_i$ を追加する。偶数回目のとき、集合 $S$ から $x_i$ を取り除く。

また、$i=1, 2, \cdots Q$ について、$i$ 回目のクエリ終了時点での $|S|$ の値は、問題文の通りにsetを管理することで得られます。これをリストにしておきます(下のコードのリスト $c$)。

ここまでのコードはこんな感じ。

N, Q = map(int, input().split())
X = list(map(lambda v: int(v) - 1, input().split()))

G = [[] for _ in range(N)]
for i, x in enumerate(X):
    G[x].append(i)

se = set()
c = [0] * Q
for i, x in enumerate(X):
    if x in se:
        se.remove(x)
    else:
        se.add(x)
    c[i] = len(se)

$j=1, 2, \cdots N$ について、クエリの中で $1$ 番目に $j$ が出てきたクエリ番号を $q_{j, 1}$、$2$ 番目に $j$ が出てきたクエリ番号を $q_{j, 2}$ とします。このとき、 $A_j$ には、$c[q_{j, 1}] + c[q_{j, 1} + 1] + \cdots + c[q_{j, 2}-1]$ を加算します。
似たような計算を何度もすることになるので、計算回数を減らすためにリスト $c$ の累積和をつくっておきます。

from itertools import accumulate

acc = [0] + list(accumulate(c))

あとは $j=1, 2, \cdots N$ について、 $S$ 内に入ったタイミングと出たタイミングがすべてわかっているので、 $j$ が $S$ 内に入っている間のリスト $c$ の値を足せばokです。

添え字がごちゃごちゃしがちなので注意しましょう。

コード

from itertools import accumulate

N, Q = map(int, input().split())
X = list(map(lambda v: int(v) - 1, input().split()))

G = [[] for _ in range(N)]
for i, x in enumerate(X):
    G[x].append(i)

se = set()
c = [0] * Q
for i, x in enumerate(X):
    if x in se:
        se.remove(x)
    else:
        se.add(x)
    c[i] = len(se)
acc = [0] + list(accumulate(c))

ans_lst = [0] * N
for ai, g in enumerate(G):
    if len(g) % 2 == 1:
        g.append(Q)  # こうしておくと、実装がラクです。
    for gi in range(0, len(g), 2):
        si, fi = g[gi], g[gi + 1]
        ans_lst[ai] += acc[fi] - acc[si]

print(*ans_lst)

さいごに

2023年5月20日のABC302以降、毎週「ABCxxxをPythonで解いてみたよ。」の記事を書き続けてきました。
いいねを付けていただいたりTwitterでも反応をくださったりして、周りの皆さんのおかげでここまで続けてこれました。
今年4月から私自身の生活が変わるので、今回でこの「ABCxxxをPythonで解いてみたよ。」の記事はいったん終了とさせていただきます。
ありがとうございました。