公式

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;
}

投稿日時:
最終更新: