AtCoder Beginner Contest 375のD問題「ABA」の解説の実装例(2)について、解説をさらにかみ砕いてメモしたいと思います。
便宜上解説のパラグラフごとに以下のように番号を振ります。
- $S_i = S_k$のとき、... ⇒ 解説の解説1
- したがって、$S_i=c$を満たす$i$を... ⇒ 解説の解説2
- これは各$X_{c,s}$ごとの寄与を考える、あるいは... ⇒ 解説の解説3
解説の解説1
$S_i = S_k$のとき、$j$は$i<j<k$を満たすものであればすべて条件を満たすので、
解説のとおり$S_i,S_j,S_k$をこの順に結合して得られる文字列が回文であることと$S_i = S_k$であることは同値なので、$S_j$は$S_i$から$S_k$の間にある文字ならどの文字を選んでもかまいません。
$j$は$k−i−1$通りの選び方があります。
たとえば3番目から6番目の文字の間の文字数は、
1 |1 2 |2 3 |3 4 |4 5 |5 6 |6
このように左から3番目の仕切りから6番目の仕切りの間にある文字数から1引いたものと考えると$6-3-1=2$となります。これを一般化して$i$番目から$k$番目の文字の間の文字数は$k−i−1$となります。
したがって、答えは$S_i = S_k$を満たす$(i,k)$の組に対する$k−i−1$の和です。
すべての$(i,k)$の組に対する$j$の選び方の総和が回文の総数となります。上で見た通り特定の$(i,k)$の組に対する$j$の選び方は$k−i−1$通りなので、$S_i = S_k$を満たす範囲内のすべての$(i,k)$の組について$k−i−1$を求め、それらをすべて足し合わせれば求めたい回文の総数となります。
解説の解説2
したがって、$S_i=c$を満たす$i$を単調増加に並べた列を$X_c$とすると、答えは$c=A,B,...,Z$に対して$\sum_{s=1}^{|X_c|}\sum_{t=s+1}^{|X_c|}(X_{c,t}-X_{c,s}-1)$を足し合わせたものです。
具体例として入力文字列$ABCACCABACBC$を例に考えます。
文字$A$について見ると、$S_i=A$を満たす$i$は$(1,4,7,9)(=X_A)$なので、$(i,k)$をこの中から選んで$k−i−1$の和を求めれば$A$について回文の総数が求まります。$(i,k)$の選び方6通りについて$k−i−1$を書き出すと以下のようになります。
\begin{array}{l}
4-1-1=2 \\ 7-1-1=5 \\ 9-1-1=7 \\ 7-4-1=2 \\ 9-4-1=4 \\ 9-7-1=1 \end{array}
最初の3つは$\sum_{t=2}^{|X_A|}(X_{A,t}-X_{A,1}-1)$とまとめることができます（$X_{A,t}$は$X_A$の$t$番目の要素、$|X_A|$は$X_A$の要素数）。
次の2つは$\sum_{t=3}^{|X_A|}(X_{A,t}-X_{A,2}-1)$となり、最後の1つは$\sum_{t=4}^{|X_A|}(X_{A,t}-X_{A,3}-1)$となります。よってさらにまとめると$\sum_{s=1}^{|X_A|}\sum_{t=s+1}^{|X_A|}(X_{A,t}-X_{A,s}-1)$となります。
これをすべての文字$c=A,B,...,Z$について一般化して、答えは$\sum_{s=1}^{|X_c|}\sum_{t=s+1}^{|X_c|}(X_{c,t}-X_{c,s}-1)$となります。
解説の解説3
解説の解説2で見た$\sum_{s=1}^{|X_c|}\sum_{t=s+1}^{|X_c|}(X_{c,t}-X_{c,s}-1)$では$O(N^2)$の計算量が必要に見えます。このままでは非常に長い文字列の場合にタイムアウトしてしまうので、計算量を削減します。
これは各$X_{c,s}$ごとの寄与を考える、あるいは累積和$\sum_{s=1}^{s'}X_{c,s}$を計算しながら求めるなどの方法で$O(|X_c|)$時間で値を求めることができます。したがってこの計算を$c=A,B,...,Z$に対して行うことにより全体として$O(N)$時間で答えを求めることができます。
ここでも具体例として入力文字列$ABCACCABACBC$を例に考えます。
文字$A$について見ると、回文の総数を得るには$(i,k)$の選び方が6通りであることから計算を6回行う必要がありました。
しかしたとえば以下の3つの場合（$k=9$）について考えると、
\begin{array}{l}
9-1-1=7 \\ 9-4-1=4 \\ 9-7-1=1 \end{array}
これらの和は$(9-1)\times 3-(1+4+7)=12$と変形できます。これは、$X_A=(1,4,7,9)$のうち$9$より前に3回登場した$A$について登場位置の合計（すなわち$X_A$の累積和）を求めておくことで、3回必要だった計算が1回で済むということを表しています。
同様にして$k=7,4$の場合も考えると、文字$A$についての回文の総数は4回($=|X_A|$回)の計算で済ませられます。
計算の回数という点で考えると、累積和を使用しない場合の計算回数は
$$ _4C_2 + _3C_2 + _5C_2 =6+3+10=19 $$
に対し、累積和を使用した場合の計算回数は
$$ 4+3+5=12 $$
となり、文字列の長さ分の計算回数で答えを出すことができます。
これを一般化します。ある文字$c$について$\sum_{s=1}^{|X_c|}\sum_{t=s+1}^{|X_c|}(X_{c,t}-X_{c,s}-1)$の各項を書き下すと以下のようになります。
\begin{array}{ccccc}
X_{c,2}-X_{c,1}-1 & & & & \\
X_{c,3}-X_{c,1}-1 & X_{c,3}-X_{c,2}-1 & & & \\
X_{c,4}-X_{c,1}-1 & X_{c,4}-X_{c,2}-1 & X_{c,4}-X_{c,3}-1 & & \\
\vdots & \vdots & \vdots & \ddots & \\
X_{c,|X_c|}-X_{c,1}-1 & X_{c,|X_c|}-X_{c,2}-1 & X_{c,|X_c|}-X_{c,3}-1 & \cdots & X_{c,|X_c|}-X_{c,|X_c|-1}-1
\end{array}
入力文字列を先頭から1文字ずつ読んで計算することを想定すると、これを横1行ごとに足したものを上から順番に足し合わせることになります。ある$t$における横1行の和
$$(X_{c,t}-X_{c,1}-1)+(X_{c,t}-X_{c,2}-1)+\cdots+(X_{c,t}-X_{c,t-1}-1)$$
は$(X_{c,t}-1)(t-1)-\sum_{s=1}^{t-1}X_{c,s}$と変形できます。$t-1$の横1行を計算する時点で$\sum_{s=1}^{t-1}X_{c,s}$（累積和）は計算できるので、$t$の横1行の和も直ちに計算できます。よって総和は$O(|X_c|)$時間で求めることができます。
$c=A,B,...,Z$について$\sum_{c}|X_c|=|S|(=N)$であることを踏まえると、同様の方法ですべての文字について和をとることで$O(N)$時間で答えを求めることができます。
コードの解説
実装例(2)のコードを以下に引用します。
#include <bits/stdc++.h>
using namespace std;
#define rep(i, n) for (ll i = 0; i < (n); i++)
using ll = long long;
int main() {
string s;
cin >> s;
vector<ll> cnt(26), sum(26);
ll ans = 0;
int n = s.size();
rep(i, n) {
int v = s[i] - 'A';
ans += (i - 1) * cnt[v] - sum[v];
cnt[v]++;
sum[v] += i;
}
cout << ans << '\n';
}
重要な変数は以下の通りです。
-
s: 入力 -
ans: 解答 -
cnt: A～Zの文字ごとの、s内での登場回数 -
sum: A～Zの文字ごとの、s内での登場位置の累積和
int v = s[i] - 'A';は、iの位置にある文字が何番目のアルファベットであるか（すなわちcntおよびsumにおける位置）を求めvに格納しています。
ans += (i - 1) * cnt[v] - sum[v];は、iの位置にある文字について、iより前までの登場回数と登場位置の累積和を用いて回文の数を計算し加算しています。