[AtCoder忘備録] ABC359をPythonで解いてみた(A～E問題) #競技プログラミング

ABC359をPythonで解いてみました。
自分なりの解き方・考え方を書いていければと思っています！！

A問題「Count Takahashi」

Diff : 10

考え方

「入力された文字列にあるTakahashiの数を数える」

入力される文字列は$N$個で、「Takahashi」か「Aoki」のどちらか
文字列を確認していき、「Takahashi」の個数だけを数える

コード

N = int(input())

ans = 0
for i in range(N):
    si = input()
    if si == "Takahashi":
        ans += 1

print(ans)

B問題「Couples」

Diff : 27

考え方

「2つ先の人が同じ服の色かを調べる」

同じ服の色を着ている人は2人だけ
$i = 0, 1, 2, ..., N - 3$番目の人と$i + 2$番目の人が同じ服を着ている場合を数える

コード

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

ans = 0
for i in range(0, 2 * N - 2):
    if A[i] == A[i + 2]:
        ans += 1

print(ans)

C問題「Tile Distance 2」

Diff : 828

考え方

「縦に1つ移動すると、横に1つ移動できると考える」

タイルは$1 × 2$の長方形で、タイル内の移動は自由にできる
これは上に1つ移動すると横にも1つ移動できることを意味している
そこで、大まかに2つの場合について考える
1. 横の移動距離$abs(T_x - S_x)$が縦の移動距離$abs(T_y - S_y)$より小さい場合
  - 縦に移動するだけで目的のタイルまでたどり着けるので、移動距離は$abs(Sy - Ty)$となる
2. 横の移動距離$abs(T_x - S_x)$が縦の移動距離$abs(T_y - S_y)$より大きい場合
  - まず目的のタイルの列まで縦移動
  - 縦に移動したときに移動できる横の移動距離を$S_x$に加減する
  - 目的のタイルの高さ$T_y$が$2$で割り切れるときは$[0, 1][2, 3][4, 5]...$とタイルが並んでいるので、$T_x$までの距離は$abs(Sx // 2 - Tx // 2)$となる
  - 目的のタイルの高さ$T_y$が$2$で割り切れないときは$[0][1, 2][3, 4]...$とタイルが並んでいるので、$T_x$までの距離は$abs((Sx + 1) // 2 - (Tx + 1) // 2)$となる

コード

Sx, Sy = map(int, input().split())
Tx, Ty = map(int, input().split())

diff = abs(Sx - Tx) - abs(Sy - Ty)
if diff > 0:
    if Sx < Tx:
        Sx += abs(Sy - Ty)
    else:
        Sx -= abs(Sy - Ty)

    # 縦の移動分
    ans = abs(Sy - Ty)
    # 横の移動分
    if Ty % 2 == 0:
        ans += abs(Sx // 2 - Tx // 2)
    else:
        ans += abs((Sx + 1) // 2 - (Tx + 1) // 2)
    
    print(ans)
else:
    print(abs(Sy - Ty))

D問題「Avoid K Palindrome」

Diff : 1381

考え方

「長さKの部分文字列でbitDPをする」

$A$、$B$、$?$で構成される文字列$S$が与えられ、$?$には$A$と$B$のどちらに置き換えてもよい
このとき、長さ$K$の部分文字列$S[0:K]$、$S[1:K+1]$、…、$S[N-1-K:N-1]$、$S[N-K:N]$について考える
- サンプル1の場合だと、$K=4$なので、$S[0:4]$、$S[1:5]$、$S[2:6]$、$S[3:7]$について考える
- $A$の場合を$0$、$B$の場合を$1$として考えると、部分文字列の文字パターンは$AAAA$～$BBBB$、つまり$0000$～$1111$となる
- 文字パターンとそれぞれの部分文字列が矛盾しない場合と、回文にならない場合を調べる($i=1$の場合、$AAAA$=$0000$はパターンとしてありえないし、$BAAB$=$0110$は回文になってしまう)
- 条件に合う文字パターンが分かったとき、$i$番目の文字パターンになるのは、$i-1$番目の共通部分列が同じ場合の組み合わせの合計になる
  - $i=1$の部分文字列が$BAAA$だったとき、いい部分文字列です
  - このとき、$i=1$番目につながる$i=0$番目の部分文字列は共通部分が$ABAA$、または$BBAA$となる
  - これより、$i=1$番目になるには、$i=0$番目で共通の要素$BAA$の先頭に$A$と$B$のつけた場合の組み合わせを足した場合になるということがわかる

     0  1  2  3  4  5  6
     A  B  ?  A  ?  B  A
i=0  ￣￣￣￣￣￣
i=1     ￣￣￣￣￣￣
i=2        ￣￣￣￣￣￣
i=3          ￣￣￣￣￣￣

? の部分には A か B のどちらかが入る

まとめると、以下のような$dp$で考えることができる

$dp[i][mask]$ ： $i$番目からの部分文字列で、$A$を$0$、$B$を$1$とする$K$ビットの$mask$になるのは何通りか
$dp[i][mask] = dp[i-1][mask >> 1] + dp[i-1][mask >> 1 | 1 << (K - 1)]$
- $dp[i-1][mask >> 1]$ ： $i-1$の先頭が$A$(=$0$)の組み合わせ
- $dp[i-1][mask >> 1 | 1 << (K - 1)]$ ： $i-1$の先頭が$B$(=$1$)の組み合わせ

コード

MOD = 998244353

N, K = map(int, input().split())
S = input()

# 文字のパターンmask と 文字列S のi文字目から始まる部分列が
# あり得る文字列、かつ回分でないか判定
def isvalid(i, mask):
    T = []
    for j in range(K):
        # S[i + j] が A で、 mask の j文字目が 1(=B) ならば不一致
        if S[i + j] == 'A' and mask >> (K - 1 - j) & 1 == 1:
            return False
        # S[i + j] が B で、 mask の j文字目が 0(=A) ならば不一致
        if S[i + j] == 'B' and mask >> (K - 1 - j) & 1 == 0:
            return False
        # S[i + j] が ? だった場合、そのまま素通り

        if (mask >> (K - 1 - j)) & 1 == 0:
            T.append('A')
        else:
            T.append('B')

    # 文字列S[i:i+K] と mask は矛盾していないが、回文かどうか判定する
    for j in range(K):
        # 回文じゃない
        if T[j] != T[K - 1 - j]:
            return True
    # 残念ながら回分です
    return False

dp = [[0] * (1 << K) for _ in range(N - K + 1)]
# 初期状態(1文字目を調べる)
for mask in range(1 << K):
    if not isvalid(0, mask):
        continue
    dp[0][mask] = 1

# 2文字目以降を調べる
for i in range(1, N + 1 - K):
    for mask in range(1 << K):
        if not isvalid(i, mask):
            continue
        a = mask >> 1
        b = mask >> 1 | 1 << (K - 1)
        dp[i][mask] += dp[i - 1][a] + dp[i - 1][b]
        dp[i][mask] %= MOD

# 答えを求める
ans = 0
for mask in range(1 << K):
    ans += dp[-1][mask]
    ans %= MOD
print(ans)

E問題「Water Tank」

Diff : 1275

考え方

「広義単調減少になる壁と水をスタックする」

それぞれの壁に水が満タンになる時間を考え、その次の時間に水があふれると考える
それぞれの壁より右側の壁が低かった場合、その区画にも水が貯まる
壁の高さ$H[i]$と、$H[i]$より低い壁の区間$w$の2つをスタックしていく
- それぞれの壁$H[i]$と低い壁の区間の長さ$w$の長方形の面積が溜まった水の量
- $H[i]$をスタックに追加するとき、スタックの最後に追加した壁の高さと比較し、壁が低かったらその分の区画を追加する区画の長さ$w$に足す
- そのとき、一旦手前の壁に溜まっている水を抜いておく
- そうすると、$H[i]$で水を貯める量が$H[i]*w$なので、計算しやすくなる
満タンになる時間の次の時間があふれる時間なので、それぞれ求めた値に$1$足した値が答えとなる

コード

from collections import deque

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

res = [0] * N
dq = deque()

# 初期条件
dq.append((0, 0))
total_water = 0

for i in range(N):
    width = 1

    # H[i] より小さい手前の壁に溜まっている水を一旦抜く
    while dq and dq[-1][0] < H[i]:
        h, w = dq.pop()
        total_water -= h * w
        width += w

    # H[i] の壁が満タンになる水の量を足す
    total_water += H[i] * width
    res[i] = total_water
    dq.append((H[i], width))

# 満水の次のタイミングがそれぞれの答え
ans = [r + 1 for r in res]
print(*ans)

まとめ・感想

　A～C問題までの三完でした。D問題とE問題は本番中に解くことができず、解説見ながら解きました。特にigasu_7さんの解説には大変お世話になりました。はじめは全然わからなくてコードを解析してなんとか…というところでした。水diff問題はやはり難しいですね。
　D問題とE問題はまだ自分の中で落とし込めていない感じがするので、考え方の書き方が雑になってしまった気がします。まだアウトプットを始めたばかりなので色々とわかりにくい解説となってしまいましたが、文章の書き方も精進していきたいところです。

[AtCoder忘備録] ABC359をPythonで解いてみた(A～E問題)

目次

A問題 「Count Takahashi」

考え方