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

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

A - Adjacent Product

問題

長さ $N$ の数列 $A=(A_1, A_2, \cdots , A_N)$ が与えられます。
$B_i = A_i \times A_{i+1}$ を満たす長さ $N-1$ の数列 $B$ を出力してください。

考察

for文をつかって、$B_1, B_2, \cdots ,B_{N-1}$ の値を順に求めていきます。
ループ回数は $N$ ではなく $N-1$ であることに注意しましょう。

リスト $B$ を出力するとき print(*B) と書くと、各要素を半角スペース区切りで出力できます。競プロでよく出てくる出力なので、知らなかった方はこの機会に覚えちゃいましょう。

コード

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

B = [0] * (N - 1)
for i in range(N - 1):
    B[i] = A[i] * A[i + 1]

print(*B)

別解

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

リスト内包表記をつかうと、そこそこシンプルに書けます。

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

B = [A[i] * A[i + 1] for i in range(N - 1)]
print(*B)

B - Piano

問題

文字列 wbwbwwbwbwbw を無限に繰り返してできる文字列を $S$ とします。
b が $B$ 個、w が $W$ 個ある区間は存在しますか？

考察

制約を見ると $W, B \leq 100$ と書いてあります。
これくらいの小さな数なら、計算量を工夫しなくてもTLEせずに済みそうです。

$i=1, 2, \cdots 12$ について、「 $S$ の $i$ 文字目から $i+W+B-1$ 文字目までを見て、 w , b がそれぞれ $W, B$ 個あるかどうか」を確認していきます。

$i$ の上限を $12$ としたのは、文字列 wbwbwwbwbwbw の長さが $12$ だからです。
$S$ は無限に wbwbwwbwbwbw のループをしているため、 $S$ の $x$ 文字目から見るのと $x+12$ 文字目から見るのは同じになります。

実際に長さ無限大の文字列をつくることはできないので、下のコードでは S = "wbwbwwbwbwbw" * 100 としています。

コード

W, B = map(int, input().split())

S = "wbwbwwbwbwbw" * 100
for i in range(12):
    T = S[i: i + W + B]
    if T.count('w') == W and T.count('b') == B:
        print("Yes")
        exit()
print("No")

別解

細かい部分をサボった解法

$W, B \leq 100$ なので、多少は計算回数が増えちゃっても大丈夫です。

S = "wbwbwwbwbwbw" * 100 としていた部分は * 1000 にしちゃってもいいですし、ループ回数を $12$ 回にしていた部分は $S$ の右端まで見てしまっても大丈夫です。

コンテスト中はACをすることが最優先なので、サボれる部分はサボってしまって次の問題にはやく取り掛かれるようにしましょう。

(下のコードは私が実際に提出したコードです。)

W, B = map(int, input().split())

S = "wbwbwwbwbwbw" * 1000
for i in range(len(S)):
    T = S[i: i + W + B]
    if T.count('w') == W and T.count('b') == B:
        print("Yes")
        exit()
print("No")

C - Σ

問題

長さ $N$ の数列 $A=(A_1, A_2, \cdots , A_N)$ が与えられます。
$1$ 以上 $K$ 以下の整数のうち、 $A$ の中に一度も現れない数の総和を求めてください。

考察1 (愚直に解いてみる)

制約を見てみると、 $1 \leq K \leq 2 \times 10^{9}$ です。

愚直に解いてみます。色々な解き方が考えられますが、例として下のコードを見てみましょう。

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

set_A = set(A)
ans = 0
for x in range(1, K + 1):
    if x not in set_A:
        ans += x

print(ans)

set内部に $x$ が含まれているかを判定する部分は $O(1)$ ですが、for文のループが $K$ 回行われていて、 $2 \times 10^9$ 回のループをしているとTLEしてしまいます。

考察2 (計算量を工夫する)

競プロあるあるの1つに、「補集合を考える」があります。

つまり...

考察1では緑色の部分を直接求めようとしていましたが、全体から灰色の部分を引いて求めることにします。

$1$ 以上 $K$ 以下の整数の総和は、 $\dfrac{K(K+1)}{2}$ です。

灰色の部分、$A$ の中に現れる $1$ 以上 $K$ 以下の整数の総和は、以下のコードで求められます。

gray = set()
for a in A:
    if 1 <= a <= K:
        gray.add(a)
sum_gray = sum(gray)

for文のループ回数も $N$ 回 ($1 \leq N \leq 2 \times 10^5$) で済んでいます。
これでこの問題を解くことができました。

コード

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

gray = set()
for a in A:
    if 1 <= a <= K:
        gray.add(a)

ans = K * (K + 1) // 2 - sum(gray)
print(ans)

D - Gomamayo Sequence

問題

0, 1 からなる長さ $N$ の文字列 $S$ が与えられます。
$i=1, 2, \cdots ,N$ について、$C_i$ 円を支払うことで $S_i$ のビットを反転させられます(0を1に、1を0にできます)。
$S$ の中で隣り合った文字が一致している箇所がちょうど1箇所にしたいです。必要な金額の最小値を求めてください。

考察

動的計画法DPで解きます。

$dp[i][v][f]$ を、左から順に $S_i$ まで見て、$S_i=v$ で、同じ数が隣り合った箇所が $f$ 個のときの、支払った金額の最小値とします。

$S_i=0$ のとき、次のような遷移をします。

• $dp[i][0][0] = dp[i-1][1][0]$
• $dp[i][0][1] = \min (dp[i - 1][0][0], dp[i - 1][1][1])$
• $dp[i][1][0] = dp[i - 1][0][0] + C[i]$
• $ dp[i][1][1] = min(dp[i - 1][0][1], dp[i - 1][1][0]) + C[i]$

$S_i=1$ のときは、上の4式の $v$ が反転するだけです。

答えは $\min (dp[N-1][0][1], dp[N-1][1][1])$ になります。

コード

INF = 1 << 60

N = int(input())
S = input()
C = list(map(int, input().split()))

dp = [[[INF] * 2 for _ in range(2)] for _ in range(N)]
if S[0] == '0':
    dp[0][0][0] = 0
    dp[0][1][0] = C[0]
else:
    dp[0][0][0] = C[0]
    dp[0][1][0] = 0
for i in range(1, N):
    if S[i] == '0':
        dp[i][0][0] = dp[i - 1][1][0]
        dp[i][0][1] = min(dp[i - 1][0][0], dp[i - 1][1][1])
        dp[i][1][0] = dp[i - 1][0][0] + C[i]
        dp[i][1][1] = min(dp[i - 1][0][1], dp[i - 1][1][0]) + C[i]
    else:
        dp[i][1][0] = dp[i - 1][0][0]
        dp[i][1][1] = min(dp[i - 1][1][0], dp[i - 1][0][1])
        dp[i][0][0] = dp[i - 1][1][0] + C[i]
        dp[i][0][1] = min(dp[i - 1][1][1], dp[i - 1][0][0]) + C[i]

ans = min(dp[N - 1][0][1], dp[N - 1][1][1])
print(ans)

E - Paint

問題

$H$ 行 $W$ 列のマス目があります。最初すべてのマス目は色 $0$ で塗られています。
$M$ 個のクエリを順に処理してください。クエリは以下の $2$ 種類です。

• 1 r x : $r$ 行目をすべて色 $x$ で塗りかえる
• 2 c x : $c$ 列目をすべて色 $x$ で塗りかえる

すべてのクエリを処理し終えたら、どの色が何マス塗られているか答えてください。

考察

競プロあるあるの1つ、「クエリを逆順に処理する」をつかいます。
今回の問題を次のように言い換えます。

$H$ 行 $W$ 列のマス目があります。最初すべてのマス目は無色で何も塗られていません。
$M$ 個のクエリが与えられます。クエリは以下の $2$ 種類です。

• 1 r x : $r$ 行目でまだ塗られていないマスをすべて色 $x$ で塗る
• 2 c x : $c$ 列目でまだ塗られていないマスをすべて色 $x$ で塗る

ただし、クエリは「逆順に」処理してください。

すべて処理し終えたら、どの色が何マス塗られているか答えてください。

どの行・どの列がクエリによって塗られたのかを管理するsetを用意して、それらを管理しながらクエリを逆順に処理することでACできます。

コード

H, W, M = map(int, input().split())
query = [tuple(map(int, input().split())) for _ in range(M)]

drawn_row, drawn_col = set(), set()
color_cnt = [0] * 200005
# 塗られてないマスは最後にすべて色0で塗る
for t, a, x in reversed(query):
    if t == 1:
        if a not in drawn_row:
            cnt = W - len(drawn_col)
            color_cnt[x] += cnt
            drawn_row.add(a)
    else:
        if a not in drawn_col:
            cnt = H - len(drawn_row)
            color_cnt[x] += cnt
            drawn_col.add(a)
color_cnt[0] = H * W - sum(color_cnt[1:])
ans_lst = []
for x in range(len(color_cnt)):
    if color_cnt[x] > 0:
        ans_lst.append((x, color_cnt[x]))

print(len(ans_lst))
for ans in ans_lst:
    print(*ans)