Official

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: