E - A > B substring Editorial by en_translator
For \(i=0,1,2,\ldots,N\), define \(A _ i\) and \(B _ i\) as the number of As and Bs in the first \(i\) characters of \(S\), respectively.
The substring formed by the \(i\)-th through \(j\)-th characters of \(S\) satisfies the condition if and only if \(A _ j-A _ {i-1}\gt B _ j-B _ {i - 1}\), or equivalently \(A _ j-B _ j\gt A _ {i-1}-B _ {i-1}\).
If we define a sequence \(D=(D _ 0,D _ 1,\ldots,D _ N)\) as \(D _ i=A _ i-B _ i\), the sought count is the number of integer pairs \((i,j)\) such that \(D _ i\lt D _ j\) and \(0\le i\lt j\le N\).
This can be evaluated in \(O(N\log N)\) time in the same manner as finding the inversion number, or in \(O(N)\) time by managing \(s _ {i,d}\coloneqq \bigl(\)the number of \(j\) with \(D _ j=d\) and \(j\leq i\)\(\bigr)\) and \(\displaystyle S _ i=\sum _ {d\lt D _ i}s _ {i,d}\).
The following is sample code.
#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] = the number of D_j so far such that D_j = d (since d ranges from -N through N, we shift the index by adding N so that it ranges from 0 through 2N.
unsigned D{N};
++counter[D];
unsigned long sum{}; // Sum of counter[d] over all d less than current D
unsigned long ans{};
for (const auto c : S) {
// Update sum and D
if (c == 'A')
sum += counter[D++];
if (c == 'B')
sum -= counter[--D];
// Update counter
++counter[D];
ans += sum;
}
cout << ans << endl;
return 0;
}
posted:
last update: