LoginSignup
1
1

ABC359のC~EをPythonで解く

Last updated at Posted at 2024-06-23

AtCoder Beginner Contest 359が最初の2問しか解けず壊滅的な結果に終わったので復習します

C - Tile Distance 2

問題文

座標平面上に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となって終わったコード

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)

解説記事を読んでやっと気づいたが、同じタイル内の場所に行くのであれば交通量は同じ(なので出発地点と到着地点をタイルの左側だとみなしてよい)ということに気づけなかった。
#通行料いらないとコメントアウトしてある箇所が最初か最後に1回しか呼ばれないことに気づければ実装できたはずなので、多分もう少しリファクタリング頑張れば時間内に解けたと思う・・・悔しい・・・

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

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

問題文

A, B, ?からなる$N$文字の文字列$S$が与えられます。
正整数$K$が与えられます。A,Bからなる文字列$T$が次の条件を満たすとき、$T$は良い文字列であるということにします。

  • $T$の長さ$K$の連続する部分文字列で、回文であるものが存在しない。

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

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

解説記事よりも解説動画の方が直感的でわかりやすかった。

この動画内で説明されているように、最悪$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文字の部分文字列は回文になってはいけないため)

\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しかないので

\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は回文にならないのでそのまま追加。

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

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

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

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

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

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

  1. $dp$を0で初期化

  2. 文字列$S$の最初の$K$文字としてありえるかつ回文ではないもの($S_K$)について$dp[0][S_K] = 1$

  3. $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] $$

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

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

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

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

ストーリー

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

問題文

長さ$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のSegTreeLazySegTreeで実装されている関数。
たとえば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$)を置いておくのが主な解決法。

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)
1
1
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
1
1