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

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

square869120Contest #2 E - 部分文字列

AtCoder 有志コン 文字列 SuffixArray prefixとsuffix lcp 総和を求める 解空間:O(N^2)通りの選択肢 区間 連続部分列を扱う問題

Suffix Array の典型問題

問題概要

英小文字のみからなる長さ N の文字列 S が与えられる。

S の連続する部分文字列として登場しうる文字列をすべて考える。重複は除外する。これらの文字列の長さの総和を求めよ。

制約

  • 1 \le N \le 10^{5}

考えたこと

次の問題とほとんど同じ！

drken1215.hatenablog.com

このリンク先の問題は、部分文字列 (重複は除外) の個数を答える問題だった。今回は「個数」ではなく「長さの総和」を答える。

とは言え、ほとんど同じ考え方で解ける。

たとえば、文字列 S = poporinri を例にとる。この文字列の Suffix Array や高さ配列 LCP を求めると次のようになる。

  • 接尾辞配列：sa = [9, 8, 5, 6, 1, 3, 0, 2, 7, 4]
  • 高さ配列：lcp = [0, 1, 0, 0, 1, 0, 2, 0, 2]

具体的には次の通りである。

0： (9 文字目以降、空文字列)
1：i (8 文字目以降)
2：inri (5 文字目以降)
3：nri (6 文字目以降)
4：oporinri (1 文字目以降)
5：orinri (3 文字目以降)
6：poporinri (0 文字目以降)
7：porinri (2 文字目以降)
8：ri (7 文字目以降)
9：rinri (4 文字目以降)

この情報をもとにして、部分文字列をすべて走査する方法を考えよう。上記の suffix たちに対して、prefix をとってできる文字列の集合から、重複を除去したものが答えとなる。

重複を除去するためには、高さ配列 lcp が役に立つ。たとえば、sa[6] を表す文字列 "poporinri" と sa[7] を表す文字列 "porinri" は、lcp 値は 2 である (下表参照)。これは、先頭から 2 文字分が重複していることを意味する。こうした重複を除去していけばよい。

Suffix Array での順序 suffix 文字列 i 番目と i+1 番目の重複分の長さ lcp[i]
0 0
1 i 1
2 inri 0
3 nri 0
4 oporinri 1
5 orinri 0
6 poporinri 2
7 porinri 0
8 ri 2
9 rinri 0

重複を除去するために、「Suffix Array の順序で後ろの方の文字列の prefix にも登場する文字列は除外して数えていく」というようにする。つまり、各 i (=0, 1, \dots, N) に対して、

  • sa[i] を表す文字列 (長さを m とする) の prefix となる文字列の長さの総和： \displaystyle \frac{m(m+1)}{2} から、 lcp[i] によって表される部分の文字列 (長さを m とする) の prefix となる文字列の長さの総和： \displaystyle \frac{l(l+1)}{2}

を引いた値を求め、その総和をとればよい。

全体として O(N) の計算量で求められる。

コード

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

// SA-IS (O(N))
template<class Str> struct SuffixArray {
    // data
    Str str;
    vector<int> sa;   // sa[i] : the starting index of the i-th smallest suffix (i = 0, 1, ..., n)
    vector<int> lcp;  // lcp[i]: the lcp of sa[i] and sa[i+1] (i = 0, 1, ..., n-1)
    int& operator [] (int i) {
        return sa[i];
    }

    // constructor
    SuffixArray(const Str& str_) : str(str_) {
        build_sa();
    }
    void init(const Str& str_) {
        str = str_;
        build_sa();
    }
    void build_sa() {
        int N = (int)str.size();
        vector<int> s;
        for (int i = 0; i < N; ++i) s.push_back(str[i] + 1);
        s.push_back(0);
        sa = sa_is(s);
        calcLCP(s);
    }

    // SA-IS
    // upper: # of characters 
    vector<int> sa_is(vector<int> &s, int upper = 256) {
        int N = (int)s.size();
        if (N == 0) return {};
        else if (N == 1) return {0};
        else if (N == 2) {
            if (s[0] < s[1]) return {0, 1};
            else return {1, 0};
        }

        vector<int> isa(N);
        vector<bool> ls(N, false);
        for (int i = N - 2; i >= 0; --i) {
            ls[i] = (s[i] == s[i + 1]) ? ls[i + 1] : (s[i] < s[i + 1]);
        }
        vector<int> sum_l(upper + 1, 0), sum_s(upper + 1, 0);
        for (int i = 0; i < N; ++i) {
            if (!ls[i]) ++sum_s[s[i]];
            else ++sum_l[s[i] + 1];
        }
        for (int i = 0; i <= upper; ++i) {
            sum_s[i] += sum_l[i];
            if (i < upper) sum_l[i + 1] += sum_s[i];
        }

        auto induce = [&](const vector<int> &lms) -> void {
            fill(isa.begin(), isa.end(), -1);
            vector<int> buf(upper + 1);
            copy(sum_s.begin(), sum_s.end(), buf.begin());
            for (auto d: lms) {
                if (d == N) continue;
                isa[buf[s[d]]++] = d;
            }
            copy(sum_l.begin(), sum_l.end(), buf.begin());
            isa[buf[s[N - 1]]++] = N - 1;
            for (int i = 0; i < N; ++i) {
                int v = isa[i];
                if (v >= 1 && !ls[v - 1]) {
                    isa[buf[s[v - 1]]++] = v - 1;
                }
            }
            copy(sum_l.begin(), sum_l.end(), buf.begin());
            for (int i = N - 1; i >= 0; --i) {
                int v = isa[i];
                if (v >= 1 && ls[v - 1]) {
                    isa[--buf[s[v - 1] + 1]] = v - 1;
                }
            }
        };
            
        vector<int> lms, lms_map(N + 1, -1);
        int M = 0;
        for (int i = 1; i < N; ++i) {
            if (!ls[i - 1] && ls[i]) {
                lms_map[i] = M++;
            }
        }
        lms.reserve(M);
        for (int i = 1; i < N; ++i) {
            if (!ls[i - 1] && ls[i]) {
                lms.push_back(i);
            }
        }
        induce(lms);

        if (M) {
            vector<int> lms2;
            lms2.reserve(isa.size());
            for (auto v: isa) {
                if (lms_map[v] != -1) lms2.push_back(v);
            }
            int rec_upper = 0;
            vector<int> rec_s(M);
            rec_s[lms_map[lms2[0]]] = 0;
            for (int i = 1; i < M; ++i) {
                int l = lms2[i - 1], r = lms2[i];
                int nl = (lms_map[l] + 1 < M) ? lms[lms_map[l] + 1] : N;
                int nr = (lms_map[r] + 1 < M) ? lms[lms_map[r] + 1] : N;
                bool same = true;
                if (nl - l != nr - r) same = false;
                else {
                    while (l < nl) {
                        if (s[l] != s[r]) break;
                        ++l, ++r;
                    }
                    if (l == N || s[l] != s[r]) same = false;
                }
                if (!same) ++rec_upper;
                rec_s[lms_map[lms2[i]]] = rec_upper;
            }
            auto rec_sa = sa_is(rec_s, rec_upper);

            vector<int> sorted_lms(M);
            for (int i = 0; i < M; ++i) {
                sorted_lms[i] = lms[rec_sa[i]];
            }
            induce(sorted_lms);
        }
        return isa;
    }

    // prepair lcp
    vector<int> rsa;
    void calcLCP(const vector<int> &s) {
        int N = (int)s.size();
        rsa.assign(N, 0), lcp.assign(N, 0);
        for (int i = 0; i < N; ++i) rsa[sa[i]] = i;
        int h = 0;
        for (int i = 0; i < N - 1; ++i) {
            int pi = sa[rsa[i] - 1];
            if (h > 0) --h;
            for (; pi + h < N && i + h < N; ++h) {
                if (s[pi + h] != s[i + h]) break;
            }
            lcp[rsa[i] - 1] = h;
        }
    }
};

int main() {
    string S;
    cin >> S;
    long long N = S.size();

    // Suffix Array の構築
    SuffixArray<string> sa(S);
    vector<int> lcp = sa.lcp;

    auto calc = [&](long long n) -> long long {
        return n * (n + 1) / 2;
    };

    // 集計していく
    long long res = 0;
    for (int i = 0; i <= N; ++i) {
        res += calc(N - sa[i]);
        if (i < N) res -= calc(lcp[i]);
    }

    // 答え
    cout << res << endl;
}