Z-algorithm や Suffix Array が使える面白い文字列検索問題！

問題へのリンク

問題概要

英小文字のみからなる長さ $N$ の文字列 $S$ が与えられる。次の条件を満たす文字列 $A, B, C$ の組が何個あるかを答えよ。

$A, B, C$ はいずれも空文字ではない
$A + B + C + A + C = S$ である

制約

$5 \le N \le 2 \times 10^{5}$

考えたこと

この手の問題では、何かをうまく固定して考えることが肝要だが......色々試行錯誤した末に、

ABC | AC

の区切れ目全探索することにした。左から $i$ 個の文字と右側の $N-i$ 個の文字とに分けることにする。このとき、 $S$ の左右反転した文字列を $T$ として、

文字列 $S$ において、S[0:] と S[i:] の先頭からの文字の一致の長さを $s$
文字列 $T$ において、T[0:] と T[N-i:] の先頭からの文字の一致の長さを $t$

とする。なお、 $s, t$ の値は、あらかじめ Z 法や Suffix Array などを用いて $O(N)$ または $O(N \log N)$ の計算量で前処理しておくことで、 $O(1)$ の計算量で求められる。

まず、 $s + t \lt N - i$ ならば、条件を満たすように文字列 $A, C$ を取ることは不可能だ。 $s + t \ge N - i$ であるとしよう。

このとき、文字列 $A$ の長さは $1, 2, \dots, \min(s, N-i-1)$ があり得て、文字列 $C$ の長さは $1, 2, \dots, \min(t, N-i-1)$ があり得る。

文字列 $A, C$ の長さの和が $N-i$ となるような長さの組を選べばよい。そのような長さの組の個数は $O(1)$ の計算量で求められる。

全体の計算量は $O(N)$ または $O(N \log N)$ となる。

コード

#include <bits/stdc++.h>
using namespace std;

vector<int> Zalgo(const string &S) {
    int N = (int)S.size();
    vector<int> res(N);
    res[0] = N;
    int i = 1, j = 0;
    while (i < N) {
        while (i+j < N && S[j] == S[i+j]) ++j;
        res[i] = j;
        if (j == 0) {++i; continue;}
        int k = 1;
        while (i+k < N && k+res[k] < j) res[i+k] = res[k], ++k;
        i += k, j -= k;
    }
    return res;
}

void ARC_055_C() {
    string S, T;
    cin >> S;
    T = S;
    reverse(T.begin(), T.end());
    int N = (int)S.size();
    
    const auto &zs = Zalgo(S);
    const auto &zt = Zalgo(T);
    
    long long res = 0;
    for (int i = N/2+1; i < N; ++i) {
        int slen = zs[i];
        int tlen = zt[N-i];
        if (slen + tlen < N-i) continue;
        
        int left = max(1, N-i-tlen);
        int right = min(N-i-1, slen);
        res += right - left + 1;
    }
    cout << res << endl;
}

int main() {
    ARC_055_C();
}

けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ARC 055 C - ABCAC (試験管黄色)

問題概要

制約

考えたこと

コード