E - A > B substring 解説
by
MMNMM
\(i=0,1,2,\ldots,N\) に対して、\(S\) の先頭 \(i\) 文字の中にある A の個数を \(A _ i\) 、B の個数を \(B _ i\) とします。
\(S\) の \(i\) 文字目から \(j\) 文字目までを取って得られる部分文字列が条件を満たすことは、\(A _ j-A _ {i-1}\gt B _ j-B _ {i - 1}\) と言い換えることができます。これは式変形をすることで \(A _ j-B _ j\gt A _ {i-1}-B _ {i-1}\) となります。
数列 \(D=(D _ 0,D _ 1,\ldots,D _ N)\) を \(D _ i=A _ i-B _ i\) で定めると、求める値は \(D _ i\lt D _ j\) かつ \(0\le i\lt j\le N\) を満たす整数の \(2\) つ組 \((i,j)\) の個数です。
これは、転倒数を求めるのと同じような要領で \(O(N\log N)\) 時間で求めたり、\(s _ {i,d}\coloneqq \bigl(D _ j=d\) かつ \(j\leq i\) を満たす \(j\) の個数\(\bigr)\) および \(\displaystyle S _ i=\sum _ {d\lt D _ i}s _ {i,d}\) の値を管理することで \(O(N)\) 時間で求めることができます。
実装例は以下のようになります。
#include <iostream>
#include <vector>
int main() {
using namespace std;
unsigned N;
string S;
cin >> N >> S;
vector<unsigned> counter(2 * N + 1); // counter[d + N] = これまでの D_j のうち D_j = d を満たすものの個数(d は -N から N の範囲をとるので、N を足して 0 から 2N にしている)
unsigned D{N};
++counter[D];
unsigned long sum{}; // 現在の D 未満の d に対する counter[d] の総和
unsigned long ans{};
for (const auto c : S) {
// sum と D を更新
if (c == 'A')
sum += counter[D++];
if (c == 'B')
sum -= counter[--D];
// counter を更新
++counter[D];
ans += sum;
}
cout << ans << endl;
return 0;
}
投稿日時:
最終更新:
