AtCoder Regular Contest 055 C - ABCAC

問題

 文字列S(|S| \le 2 \times 10^5)が与えられるので、それを空でない文字列A, B, Cを用いてABCAC (文字列を連結したもの)と分割する方法が何通りあるかを求めよ。

解法

 ABC / ACと分ける切れ目を全探索する。分けられた前半部分と後半部分で、prefixとsuffixがそれぞれ何文字一致しているかが分かれば、それをA,Cとして分割する方法が何通りあるかはO(1)で求められる。

 このprefix/suffixが何文字一致しているかはあらかじめもとのSに対してZ-algorithmを適用すれば先頭に対する最長共通接頭辞の長さがO(|S|)でわかる。Sを反転したものに対してもZ-algorithmを適用することで末尾に対する最長共通接尾辞の長さもO(|S|)でわかるため、この問題が解ける。

反省

 2時間10分(1WA)でAC。1時間ほど考えてわからなかったのでとりあえず嘘っぽい部分点解法を出して諦めた。部分点解法は2回目のAが出てくる位置、及びAの長さをループし、さらに文字列比較でその長さ分だけかかるので、O(|S|^3)になっているのではないかという気がするが、800msecほどでなんとか部分点は取れだ

 解説を読むと、最初はKMPとかを使って解くのかと思ったんだけど、書いてあるようにZ-algorithmを使えばスマートらしい。Z-algorithmはよく知らなかったので多少調べて、なんとなく確かに上手くいきそうだなーくらいの理解で他人の実装を写して出した。答えとして調整する部分もいまいち理解が浅くてもう一度出されても苦労しそうな問題に思える。

コード

#include"bits/stdc++.h"
using namespace std;
using ll = int64_t;

//A[i] := SとSのi文字目以降の最長共通接頭辞の長さ
//としたものを返す
//参考:http://snuke.hatenablog.com/entry/2014/12/03/214243
vector<ll> Z_algorithm(string S) {
    vector<ll> A(S.size());
    A[0] = S.size();
    ll i = 1, j = 0;
    while (i < S.size()) {
        while (i + j < S.size() && S[j] == S[i + j]) {
            j++;
        }
        A[i] = j;
        if (j == 0) {
            i++;
            continue;
        }
        ll k = 1;
        while (i + k < S.size() && k + A[k] < j) {
            A[i + k] = A[k];
            k++;
        }
        i += k;
        j -= k;
    }
    return A;
}

int main() {
    string S;
    cin >> S;

    vector<ll> z = Z_algorithm(S);
    reverse(S.begin(), S.end());
    vector<ll> rz = Z_algorithm(S);
    reverse(rz.begin(), rz.end());

    const ll n = S.size();
    ll ans = 0;

    for (ll i = n / 2 + 1; i < n; i++) {
        ll a = min(n - i - 1, z[i]);
        ll c = min(n - i - 1, rz[i - 1]);
        ans += max(min(a + c - (n - i) + 1, n - i - 1), (ll)0);
    }

    cout << ans << endl;
}