AtCoder ABC214 F - Substrings の 公式解説 は、まず「印をつけた文字が隣り合ってはいけない」という制約を外した問題から解説されている。その問題は次のとおりである。
$S$ の空文字列でない部分文字列の個数を求めよ。ただし、部分文字列とは元の文字列から $0$ 文字以上取り除き残った文字を元の順番で並べたものである。
そしてその解説の中で、計算量に関して次の様に書かれている。
計算量は一見 $O(|S|^2)$ に見えますが、文字の種類数を $σ$ としたとき、累積和を用いれば $O(|S|)$、用いなくとも $O(σ|S|)$ で十分高速です。
そこで、累積和を用いて計算量が $O(|S|)$ であるコードに挑戦してみた。
累積和を ${\rm{dp}}_i$ で表したいので、公式解説の ${\rm{dp}}_i$ を $a_i$ に変えた次のような数列 {$a_i$} を考えます。
- $a_i$ := 文字列 $S$ の $1$ 文字目から $i$ 文字目までの部分から得られる部分文字列のうち $i$ 文字目は必ず使うようなものの個数
ただし、空文字列に対応して $a_0=1$ とします。
一般項 $a_i$ は次のようになります。（理由は公式解説 を見てください。）
- $k$ を $S_i=S_k$ かつ $k<i$ を満たす最大の整数として（存在しない場合 $k=0$）
a_i=\sum_{j=k}^{i-1}a_j
ここまで、引用ばかりですが、次から自分で考えた部分です。
${\rm{dp}}_i$ を数列 {$a_i$} の和とします。
{\rm{dp}}_i=\sum_{j=0}^{i}a_j
このとき、遷移式は次の様になります。
\begin{eqnarray}
{\rm{dp}}_{i+1} &=& \sum_{j=0}^{i+1}a_j \\
&=& a_{i+1}+\sum_{j=0}^{i}a_j \\
&=& \sum_{j=k}^{i}a_j+{\rm{dp}}_i \\
&=& \sum_{j=0}^{i}a_j-\sum_{j=0}^{k-1}a_j+{\rm{dp}}_i \\
&=& {\rm{dp}}_i-{\rm{dp}}_{k-1}+{\rm{dp}}_i \\
&=& 2{\rm{dp}}_i-{\rm{dp}}_{k-1}
\end{eqnarray}
ただし、$k=0$ に対する ${\rm{dp}}_{-1}$ は
{\rm{dp}}_{-1}=0
とする。
これをpythonのコードにした。
substrings.py
S = input()
N = len(S)
# k は k<i で S[k]=S[i] となる最大の index (存在しないときは -1)
k = [-1] * 26
dp = [0] * (N + 2) # k=-1 のために N+2 とした（dp[-1]=dp[n+1]=0）
dp[0] = 1 # 初期値のため空文字列 '' を 1 個と数える
for i in range(N):
dp[i + 1] = 2 * dp[i] - dp[k[ord(S[i]) - ord('a')]]
k[ord(S[i]) - ord('a')] = i # S[i] に対する k を最大値に更新
print(dp[N] - 1) # 空文字列は部分列に数えないので dp[0]=1 を引く
# S='abaaabba' の実行結果 59
計算量を $O(|S|)$ に出来たと思うが、どなたかアドバイスをお願いします。