ABC104 - kamojirobrothersのブログ

ABC104のD問題

ABC104のD問題で'B'を固定する解法についての解説。
D: We Love ABC - AtCoder Beginner Contest 104 | AtCoder
問題は次のようなものです。
ABCからなる文字列Tに対し、i < j < k番目の文字がそれぞれA, B, Cになるような(i, j, k)の組み合わせの数をABC数と呼びます。
ABC?からなる文字列Sが与えられたときに、?はABCそれぞれが入り、 ${3^Q}$ 通りあります。
各場合に関するABC数の和を求め、それを ${10^9+7}$ で割った余りを求める問題です。

解説

i番目のBを固定して、左にAがある場合の数とCがある場合の数を求めます。
Aの場合の数についてのみ解説します。
Aの数をn、?の数をmとします。
すでにあるAからAを選ぶ場合の数は、Aの数nと?に何を入れるかの ${3^m}$ の積です。
次に?からAを選ぶ場合は、?のうちのk個がAである場合それぞれについて考えます。
?にAがk個入る場合の数は ${ _nC_k}$ 通りあり、その中のどのAを選ぶかはk通りあります。
また、残りの ${(n-k)}$ 個の?にB,Cのいずれが入るかの ${2^{n-k}}$ 通りあります。
これらをまとめると
${n \times 3^m + \sum_{k=1}^n \ _nC_k \times k \times 2^{n-k}}$
となります。
第2項のΣは、 ${(2+x)^n}$ を微分することで簡単にできます。
というわけで、
${n\times 3^m + n \times 3^{n-1}}$
がAの場合の数となります。
Aの場合の数とCの場合の数の積がi番目にBがあるときのABC数の合計になります。
プログラムは次のものになります。

S = input()
L = len(S)
Q = 10**9 + 7
anum = [ 0 for _ in range(L+1)] #i番目までのAの個数の合計
cnum = [ 0 for _ in range(L+1)] #i番目までのCの個数の合計
hnum = [ 0 for _ in range(L+1)] #i番目までの?の個数の合計
for i in range(1,L+1):
    if S[i-1] == 'A':
        anum[i] = anum[i-1] + 1
        cnum[i] = cnum[i-1]
        hnum[i] = hnum[i-1]
    elif S[i-1] == 'B':
        anum[i] = anum[i-1]
        cnum[i] = cnum[i-1]
        hnum[i] = hnum[i-1]
    elif S[i-1] == 'C':
        anum[i] = anum[i-1]
        cnum[i] = cnum[i-1] + 1
        hnum[i] = hnum[i-1]
    else:
        anum[i] = anum[i-1]
        cnum[i] = cnum[i-1]
        hnum[i] = hnum[i-1] + 1
ans = 0
czen = cnum[L] #全部のCの個数
hzen = hnum[L] #全部の?の個数
for i in range(1,L+1):
    if S[i-1] == 'B' or S[i-1] == '?':
        A = ((anum[i-1]*pow(3,hnum[i-1],Q))%Q + (hnum[i-1]*pow(3,max(0,hnum[i-1]-1),Q))%Q)%Q
        C = (((czen - cnum[i])*pow(3,hzen - hnum[i],Q))%Q + ((hzen - hnum[i])*pow(3,max(0,hzen - hnum[i]-1),Q))%Q)%Q
        K = (A*C)%Q
        ans = (ans + K)%Q
print(int(ans))

おまけ

powではなく**でやってたときは、TLEになった。