ABC359のC~EをPythonで解く

[AtCoder Beginner Contest 359](https://atcoder.jp/contests/abc359)が最初の2問しか解けず壊滅的な結果に終わったので復習します

# [C - Tile Distance 2](https://atcoder.jp/contests/abc359/tasks/abc359_c)

## 問題文
> 座標平面上に2×1 の大きさのタイルが敷き詰められています。 タイルは、次の規則に従って敷き詰められています。整数の組$(i,j)$に対し、正方形$A_{i,j}={(x,y)∣i≤x≤i+1∧j≤y≤j+1}$は1つのタイルに含まれる。

> i+j が偶数のとき、$A_{i,j}$と$A_{i+1,j}$は同じタイルに含まれる。ただし、タイルは境界を含むものとし、共通部分が正の面積をもつような2つの異なるタイルは存在しないとします。

> 高橋君は、はじめ座標平面上の点$(S_x+0.5,S_y+0.5)$にいます。高橋君は、次の移動を好きなだけ繰り返します。
    - 上下左右の方向と正の整数$n$を選ぶ。その方向に$n$だけ進む。
高橋君が異なるタイルを通るたび、高橋君は通行料を1だけ支払います。高橋君が点$(T_x+0.5,T_y+0.5)$にたどり着くために支払わなければならない通行料の最小値を求めてください。

以下、コンテスト時間中に提出してTLEとなって終わったコード

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

x,y = Sx,Sy
count = 0
while x!=Tx or y!=Ty:
    # 左に行きたいとき
    if x > Tx:
        if (x+y)%2==1:
            x -= 1 # 通行料いらない
        elif y > Ty:
            y -= 1
            x -= 1
            count += 1
        elif y < Ty:
            y += 1
            x -= 1
            count += 1
        else:
            x -= 1
            count += 1
    # 右に行きたいとき
    elif x < Tx:
        if (x+y)%2==0:
            x += 1 # 通行料いらない
        elif y > Ty:
            y -= 1
            x += 1
            count += 1
        elif y < Ty:
            y += 1
            x += 1
            count += 1
        else:
            x += 1
            count += 1
    # あとは上下動くだけのとき
    else:
        if y > Ty:
            y -= 1
            count += 1
        else:
            y += 1
            count += 1
    #print(f"x={x},y={y}")
print(count)
```
[解説記事](https://atcoder.jp/contests/abc359/editorial/10264)を読んでやっと気づいたが、**同じタイル内の場所に行くのであれば交通量は同じ**(なので出発地点と到着地点をタイルの左側だとみなしてよい)ということに気づけなかった。
`#通行料いらない`とコメントアウトしてある箇所が最初か最後に1回しか呼ばれないことに気づければ実装できたはずなので、多分もう少しリファクタリング頑張れば時間内に解けたと思う・・・悔しい・・・

以下、解説記事を参考に書き換えた記事。これなら最大でも2ループで終了する。

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

x,y = Sx,Sy
count = 0

if (x+y)%2==1:
    x-=1
if (Tx+Ty)%2==1:
    Tx-=1
while x!=Tx or y!=Ty:
    # 左に行きたいとき
    m = min(abs(x-Tx),abs(y-Ty))
    if x > Tx:
        if y > Ty:
            y -= m
            x -= m
            count += m
        elif y < Ty:
            y += m
            x -= m
            count += m
        else:
            m = (x-Tx)//2
            x -= m*2
            count += m
    # 右に行きたいとき
    elif x < Tx:
        if y > Ty:
            y -= m
            x += m
            count += m
        elif y < Ty:
            y += m
            x += m
            count += m
        else:
            m = (Tx-x)//2
            x += 2*m
            count += m
    # あとは上下動くだけのとき
    else:
        if y > Ty:
            count += y-Ty
            y = Ty
        else:
            count += Ty-y
            y=Ty
    #print(f"x={x},y={y}")
print(count)
```

# [D - Avoid K Palindrome](https://atcoder.jp/contests/abc359/tasks/abc359_d)

## 問題文

> `A`, `B`, `?`からなる$N$文字の文字列$S$が与えられます。
正整数$K$が与えられます。`A`,`B`からなる文字列$T$が次の条件を満たすとき、$T$は良い文字列であるということにします。
> - $T$の長さ$K$の連続する部分文字列で、回文であるものが存在しない。

> $S$に含まれる`?`の個数を$q$個とします。$q$個の`?`をそれぞれ`A`,`B`のどちらかに置き換えて得られる文字列は$2^q$通りありますが、その中に良い文字列がいくつあるか求めてください。ただし、答えは非常に大きくなる場合があるので、998244353 で割った余りを求めてください。

**これは解き方すら思いつかず飛ばした**
動的計画法苦手かもしれない。復習しなくては。

[解説記事](https://atcoder.jp/contests/abc359/editorial/10265)よりも[解説動画](https://www.youtube.com/watch?v=FGdRakXIEXM&feature=youtu.be)の方が直感的でわかりやすかった。

<iframe width="560" height="315" src="https://www.youtube.com/embed/FGdRakXIEXM?si=PpnET5DWSNeiV_xB&amp;start=108" title="YouTube video player" frameborder="0" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture; web-share" referrerpolicy="strict-origin-when-cross-origin" allowfullscreen></iframe>

この動画内で説明されているように、最悪$2^q (q=文字列Sに含まれている?の数)$だけ全探索すればよいが、この問題では$?$の数に制約がついていないため最大$2^{1000}$通り調べなければならなくなる。

そこで、文字列$S$の始めの$K$文字だけをまず考えて、その結果に応じて$1文字目~K+1文字目$,$2文字目~K+2文字目$も考えていく**動的計画法**の手法をとる。

例えば入力例1を見てみる
```
7 4
AB?A?BA
```

最初の4文字`AB?A`だけを見ると、3文字目の`?`に当てはまるのは`A`しかない。(4文字の部分文字列は回文になってはいけないため)

```math
\begin{align}
    dp[0][ABAA] &= 1\\
    dp[0][\{ABAA以外の4文字の文字列\}] &= 0
\end{align}
```

※以下、便宜上文字の数え方がプログラムのインデックスに従っている(0文字目~)

次に1~4文字目を見るが、その前に$dp[0][S_4]>0$となっているような文字列$S_4$を確認すると`ABAA`しかないので、1~3文字目に当てはまるのはその最後の3文字`BAA`しかない。

$S$の4文字目は?だが、`BAA`の末尾につけて回文にならないのは`A`しかないので
```math
\begin{align}
    dp[1][BAAA] &= dp[1][BAAA] + dp[0][ABAA]\\
    dp[1][\{BAAA以外の4文字の文字列\}] &= 0
\end{align}
```

同様に2~5文字目を見るとき、2~4文字目として当てはまるのは`AAA`しかない。5文字目は`B`であり`AAAB`は回文にならないのでそのまま追加。

```math
\begin{align}
    dp[2][AAAB] &= dp[2][AAAB] + dp[1][BAAA]\\
    dp[2][\{AAAB以外の4文字の文字列\}] &= 0
\end{align}
```

最後に3~6文字目も同様に計算する

```math
\begin{align}
    dp[3][AABA] &= dp[2][AABA] + dp[2][AAAB]\\
    dp[3][\{AABA以外の4文字の文字列\}] &= 0
\end{align}
```

すると、以下のように答えを求めることができる

```math
answer = \sum_{S_4 \in \{AとBからなる4文字の文字列\}} dp[3][S_4]
```

これを一般化すると以下のようになる。

0. $dp$を0で初期化
1. 文字列$S$の最初の$K$文字としてありえるかつ回文ではないもの($S_K$)について$dp[0][S_K] = 1$
2. $i \geq 1$のとき、$i$文字目~$i+K-2$文字目としてあり得るもの(つまり$dp[i-1][S_K] > 0$である$S_K[1:]$)を取ってくる。それに$i+K-1$文字目(`?`であれば$i+K-1$文字目としてありえるもの)を末尾に付け加えたもののうち回文になっていない文字列$\hat{S}_K$に対して
$$dp[i][\hat{S}_K] = dp[i][\hat{S}_K] + dp[i-1][S_K] $$

3. 最終的に総和を求める

```math
answer = \sum_{S_K \in \{AとBからなるK文字の文字列\}} dp[N-K][S_K]
```

また、この問題では$998244353$で割った余りを出力しなければならない点に注意しながら実装すると以下のようになる。今回文字列をビットに置き換えて実装したが、辞書形式などを使って文字列をそのままキーとして扱うのも悪くない。

```python
from copy import copy

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

# n文字目~n+K-1文字目がK桁の2進数に対応する文字かつ~n+K-1文字目以前が条件を満たすものの個数
matrix = [[0] * (1 << K) for _ in range(N-K+1)]

def kaibun_hantei(s):
    if len(s) == 1:
        return True
    elif len(s) == 2:
        return s[0] == s[1]
    elif s[0] == s[-1]:
        return kaibun_hantei(s[1:-1])
    else:
        return False
    
def bit_to_str(n):
    c = []
    for i in range(K):
        if (n >> i) & 1:
            c.append('A')
        else:
            c.append('B')
    return c

bit_to_str_dict = {n:bit_to_str(n) for n in range(1 << K)}
str_to_bit_dict = {''.join(bit_to_str(n)):n for n in range(1 << K)}
kaibun_hantei_bit_dict = {n:kaibun_hantei(bit_to_str_dict[n]) for n in range(1 << K)}

def predict_bits_from_str(s):
    hatena_indices = [i for i in range(len(s)) if s[i]=='?']
    bits = []
    for i in range(1 << len(hatena_indices)):
        s_tmp = s.copy()
        for j in range(len(hatena_indices)):
            if (i >> j) & 1:
                s_tmp[hatena_indices[j]] = 'A'
            else:
                s_tmp[hatena_indices[j]] = 'B'
        s_tmp = ''.join(s_tmp)
        bit = str_to_bit_dict[s_tmp]
        if kaibun_hantei_bit_dict[bit] == False:
            bits.append(bit)
    #print(bits)
    return bits

for bit in predict_bits_from_str(S[:K]):
    matrix[0][bit] = 1
    
if N > K:
    for i in range(1,N-K+1):
        for bit in predict_bits_from_str(S[i-1:i+K-1]):
            if matrix[i-1][bit] == 0:
                continue
            s = bit_to_str_dict[bit]
            #print(s,"->",s[1:]+[S[i+K-1]])
            for bit_new in predict_bits_from_str(s[1:]+[S[i+K-1]]):
                matrix[i][bit_new] = (matrix[i][bit_new] + matrix[i-1][bit]) % MOD
            
count = 0
for j in range(1 << K):
    count = (count + matrix[N-K][j]) % MOD
print(count)
```

# [E - Water Tank](https://atcoder.jp/contests/abc359/tasks/abc359_e)

## ストーリー
> 長い水槽があり、高さの異なる板が等間隔に配置されています。 高橋くんは、この水槽の端へ水を注いでいったとき、板で区切られたそれぞれの領域に水が到達する時刻が知りたいです。

## 問題文
> 長さ$N$の正整数列$H=(H_1,H_2,…,H_N)$が与えられます。長さ$N+1$の非負整数列$A=(A_0,A_1,…,A_N
)$があります。 はじめ、$A_0=A_1=⋯=A_N=0$です。

> Aに対して、次の操作を繰り返します。
    1. $A_0$の値を1増やす。
    2. $i=1,2,…,N$に対して、この順に次の操作を行う。
        - $A_{i−1}>A_i$かつ$A_{i−1}>H_i$のとき、$A_{i−1}$の値を1減らし、$A_i$の値を1増やす。
        - $i=1,2,…,N$のそれぞれに対して、初めて$A_i>0$が成り立つのは何回目の操作の後か求めてください。


セグメント木やら遅延セグメント木やらを使いそうだけど`min_left`ってなんだっけ覚えてない(泣)と焦っていたらコンテスト時間終了してました。

でも後でAtCoder Libraryのソースコード読みながらやったらできたので、だいたいの理論は間違っていなかったっぽい。

おおまかな理論は以下

1. 水が0個目の板を超えるのは明らかに時刻`H[0]+1`のとき。このとき0個目の板の手前には`H[0]`の水が溜まっている。
2. $i \geq 1$のとき、$i$個目の板の高さ以上の板が$i-1$枚目までにあるかどうかを確認する。
    a. もしあれば、そのような板のうち一番$i$枚目に近い場所に立っている板の場所を$j$とする。$i$個目の板を水が超えるとき$j$枚目までの板の手前にある水はそのままなので、$j+1~i$を`H[i]`ぶんの水でいっぱいにする。このとき入っている水の総量+1が水が$i$番目の板を超える時刻
    b. なければ$i$枚目までの板の手前にある水を`H[i]`でいっぱいにして、このとき入っている水の総量+1が水が$i$番目の板を超える時刻

2.の部分を実装するときに`min_left`が必要になってくる

## min_leftとは？

AtCoder Libraryの`SegTree`や`LazySegTree`で実装されている関数。
たとえば`tree`をセグメント木もしくは遅延セグメント木、`func`をセグメント木や遅延セグメント木の要素に対して処理を行う関数であれば、`tree.min_left(r,func)`を実行すると`func(tree.prod(l,r))`を行った際に`True`が返される最大の`l`を返してくれる。

同じような関数である`tree.max_right(l,func)`は`func(tree.prod(l,r))`を行った際に`True`が返される最小の`r`を返してくれる。


関数の名前がいろいろとややこしいが覚えるしかない



入力例の`H`をそのまま`min_left`すると、返り値が`0`だったときにa.のケースにおける求める板が0枚目なのか、それともb.のケースなのかが分かりづらいので注意が必要。今回の場合だと0の時に改めてチェックし直すか、0枚目の前に十分に高い板($10^9+1$)を置いておくのが主な解決法。

```python
from atcoder.lazysegtree import LazySegTree
from atcoder.segtree import SegTree

N = int(input())
H = list(map(int,input().split()))
H_max_left = [0] * N

def op(data1,data2):
    v,l = data1
    vv,ll = data2
    #print(data1,data2,(v+vv,l+ll))
    return (v+vv,l+ll)
e = (0,0)
def mapping(up_lazy,sum_data):
    v,l = sum_data
    return (max(v,l*up_lazy), l)

def composition(up_lazy,down_lazy):
    return max(up_lazy,down_lazy)

id_ = -float("inf")

lst = LazySegTree(op,e,mapping,composition,id_,[(0,1)]*N)
lst2 = SegTree(max,-1,H)
    
answer = []
for i in range(N):
    #print(W)
    #print(answer)
    #ind = binary_search(0,i-1,H[i])
    ind = lst2.min_left(i,lambda a:a<H[i])
    #print(ind)
    #print(f"ind={ind},{H[i]}")
    if ind==0:
        if H[0] >= H[i] and i > 0:
            lst.apply(ind,i+1,H[i])
            r = lst.prod(0,i+1)
            answer.append(r[0]+1)
        else:
            lst.apply(0,i+1,H[i])
            r = lst.prod(0,i+1)
            answer.append(r[0]+1)
    else:
        lst.apply(ind,i+1,H[i])
        r = lst.prod(0,i+1)
        answer.append(r[0]+1)
    #print([lst.get(i)[0] for i in range(N)])
print(*answer)
```