AtCoder Beginner Contest 104 - かつかれーのメモ帳

真面目にやり始めて日が浅いので、得た学びは遡り切れるだろうと考えて遡る。忙しくはないけど院試などがあり、微妙に時間がなくて悲しい。これも出られていないけどDだけ解いた。
コンテストページ: AtCoder Beginner Contest 104 - AtCoder

D: We Love ABC

'A','B','C','?'(ワイルドカード)のみから成る文字列が与えられる。
'?'をそれぞれ好きに'A','B','C'のどれかに置き換えると、'?'の個数を $Q$ として $3^Q$ 通りの文字列ができる。
それぞれの文字列 $T$ について部分列(連続していなくても良い)*1として"ABC"が取れる箇所を数え、その個数を足しあげよ。

という問題。
入力例1の"A??C"がわかりやすいが、 $T$ が"ABAC","ABBC","ABCC"のそれぞれのパターンにおいて、0,1,3文字目を取ったものがカウントされている。これを重複してカウントしないものだとしばらく誤読していた。Ratedならまあまあ悲しくなっていたと思うので、入力例をよく見る癖をつけるべき。

ただ、実はこの誤読した問題、すなわち

$T$ の部分列(連続していなくても良い)のうち、"ABC"となりうる箇所は何箇所あるか？

が解ければあとはそんなに難しくない。それぞれの「"ABC"となりうる箇所」は、 $T$ の残り $|T|-3$ 文字に'?'が $R$ 個含まれているとすると $3^R$ 回重複して数えられており*2、この重みをつけて足し上げて行けばよいためである。'?'は全部で $Q$ 個なのだから、「"ABC"となりうる箇所」に含まれる'?'の数を $S (0 \le S \le 3)$ としたとき $R = Q - S$ が成立し、この重み付けの係数は簡単に求まる。

「"ABC"となりうる箇所」として有り得る"???","??C","?B?","A??","?BC","A?C","AB?","ABC"の8通りについて $T$ から3文字選ぶような全探索を書くと $O(|T|^3)$ で無理だが、後ろから"??","?C","B?","BC"の個数を数えていき、'A'または'?'に出会ったときにこの4通りについて処理をすれば $O(|T|)$ 。ひとまずこの方針で、以下に示すようなイケてないコードが出来上がり、通すことはできる。これではちょっとあんまりなので、よりまともな解法を検討していく。
Submission #3086889 - AtCoder Beginner Contest 104

mod = 1000000007
s = input()
n = len(s)
c = [0 for i in range(n)] #c[i]=(s[i:]に存在する'C'の個数)
d = [0 for i in range(n)] #d[i]=(s[i:]に存在する'?'の個数) 以降も、'd'は'?'に対応する
bc = [0 for i in range(n)] #bc[i]=(s[i:]に存在する"BC"の個数)
bd = [0 for i in range(n)] #上と同様、"B?"
dc = [0 for i in range(n)] #"?C"
dd = [0 for i in range(n)] #"??"


if s[-1] == "C":
    c[-1] = 1
if s[-1] == "?":
    d[-1] = 1
for i in range(-2, -n-1, -1): #後ろから2文字目～先頭を見る
    c[i] = c[i+1]
    d[i] = d[i+1]
    bc[i] = bc[i+1]
    bd[i] = bd[i+1]
    dc[i] = dc[i+1]
    dd[i] = dd[i+1]
    if s[i] == "B":
        bc[i] += c[i+1] #今見つけた'B'と、今までに見つけた'C'で"BC"をつくる。後に続くのも同様。
        bd[i] += d[i+1]
    elif s[i] == "C":
        c[i] += 1
    elif s[i] == "?":
        d[i] += 1
        dc[i] += c[i+1]
        dd[i] += d[i+1]

# d[0]は'?'の総数で、問題文中のQに対応する
if d[0] == 0: #d[0]が3に満たないときはバグるので別扱い
    mul = [0, 0, 0, 1]
elif d[0] == 1:
    mul = [0, 0, 1, 3]
elif d[0] == 2:
    mul = [0, 1, 3, 9]
else:
    mul = [1, 1, 1, 1] #以下の処理でmul=[3**(Q-3),3**(Q-2),3**(Q-1),3**Q]となる
    for i in range(d[0]-3): #3をQ-3乗
        mul[0] *= 3
        mul[0] %= mod
    for i in range(3): #あとは順に3倍
        mul[i+1] = (mul[i]*3) % mod

res = 0
for i in range(n-2): #一文字ずつ見ていって足し上げる
    if s[i] == "A":
        res += bc[i+1]*mul[3]+(dc[i+1]+bd[i+1])*mul[2]+dd[i+1]*mul[1]
    if s[i] == "?":
        res += bc[i+1]*mul[2]+(dc[i+1]+bd[i+1])*mul[1]+dd[i+1]*mul[0]
    res %= mod
print(res)

こちらhttps://www.hamayanhamayan.com/entry/2018/08/05/232127で取り上げられているように、最後の数え上げはもうちょっと楽に書ける定石があるらしい。 $1\le i \le |T|$ として、前から $i$ 文字目までの'A'の個数 $a$ と'?'の個数 $l$ 、後ろから $i$ 文字目までの'C'の個数 $c$ と'?'の個数 $r$ を初めに数えて持っておくと、改めて前から見ていって $T$ の $j$ 文字目が'B'のときに $3^Qa_{j-1}c_{j+1}+3^{Q-1}(l_{j-1}c_{j+1}+a_{j-1}r_{j+1})+3^{Q-2}l_{j-1}r_{j+1}$ 個をカウントしていけば良いことになる( $k$ 文字目が'?'のときも同様だが、ワイルドカードが一個減るので指数を1つずつ減らす)。計算量は変わらないものの、これで素直に探索を書くことができるようになる。半分全列挙とかでもそうだが、端からではなく真ん中を決めてみると楽になることが結構ありそうなので、この感覚は身につけておきたい。

想定解はDP。全探索書かないならDPだろうなとは思うけど、解説の通りのスッキリしたDPテーブルを書ける感じがしないので、これくらいの問題で練習しておくと良い練習になりそう。勉強のために、ほぼ引き写しになるが書いておく。文字列を前後反転して処理すると簡単なので、想定解とはちょっとだけ表現やDPテーブルの対応が異なる。
$dp_{i,j} (0 \le i \le |T|, 0 \le j \le 3)$ は「 $T$ の後ろから $i$ 文字目までの中から、"ABC"の後ろ $j$ 文字を選ぶ場合の数」とする。求めるのは $dp_{|T|,3}$ 。 $T$ の前後を反転させた文字列を改めて $T$ として取り直すと、 $dp_{i,j}$ が「 $T$ の前から $i$ 文字目までの中から、"CBA"の前 $j$ 文字を選ぶ場合の数」となり少し書きやすい。この場合も $dp_{|T|,3}$ を求めればよい。
推移を考えていく。特に $i+1$ 文字目が'?'のときは

$dp_{i+1,0}=3dp_{i,0}$
$dp_{i+1,1}=3dp_{i,1}+dp_{i,0}$
$dp_{i+1,2}=3dp_{i,2}+dp_{i,1}$
$dp_{i+1,3}=3dp_{i,3}+dp_{i,2}$

となる(後ろ三行の右辺にある二項のうち、後者は'?'を"ABC"の列を伸ばすために使っているので1通りしか使い方がない一方、前者は何でもよい文字として使っているので3通り使い方がある)。
その他の場合、 $i+1$ 文字目を無視すれば $dp_{i+1,j}$ は $dp_{i,j}$ と等しくなるのでとりあえず前の値をコピーする。コピー後に $i+1$ 文字目が何だったかに着目して増分を足してやれば良くて、

'C'のとき、 $dp_{i+1,1}+=dp_{i,0}$
'B'のとき、 $dp_{i+1,2}+=dp_{i,1}$
'A'のとき、 $dp_{i+1,3}+=dp_{i,2}$

となる。
やっていることは意外に全探索と変わりがないが、見通しが良いしインデックスがテーブルの推移にそのまま使えるので以下のようにコードが見やすくなる。
Submission #3086882 - AtCoder Beginner Contest 104

mod = 1000000007
s = input()[-1::-1] #入力のタイミングで文字列を反転してしまう
n = len(s)
dp = [[0]*4 for i in range(n+1)]
dp[0][0] = 1 #0文字から0文字選ぶ方法だけ1で初期化。他は0。
for i in range(n):
    for j in range(4):
        dp[i+1][j] = dp[i][j] % mod #とりあえず一つ前の値はコピー
    if s[i] == "?":
        #dp_{i+1,0}だけは一項少ないので別扱い。
        dp[i+1][0] += (2*dp[i][0]) % mod #コピーしたものに2倍を足せば合計3倍。
        for j in range(1, 4):
            dp[i+1][j] += (2*dp[i][j]+dp[i][j-1]) % mod
    else:
        index = "CBA".find(s[i])+1 #indexを1-originにするために1を加えている
        dp[i+1][index] += dp[i][index-1] % mod
print(dp[n][3] % mod)

所感

前半のアホみたいな解法はまあまずコーナーで落ちないのでその点ではよく、こどふぉやってると上手になってくる感じがある。
とはいえあんな解法ではアレなので(語彙力)、DP上手になりたい。

*1:数列だとこれを部分列と呼んでいる例を見たことあるけど、文字列でも言うのかな・・・"部分文字列"は連続していなければならない気がなんとなくする。

*2:この"ABC"3文字分だけ確定すると、残りの文字列の作り方は何通りあるか？という考え方。作りうる全ての文字列においてこの"ABC"はABC数に寄与している。