これを機会に SA-IS を整備した！　今回の記事はあくまで自分が読んでわかる以上を目指さない備忘録として。

問題へのリンク

問題概要

2 つの文字列 $X, Y$ に対して「先頭何文字が一致しているか」を $f(X, Y)$ と表すことにします。

長さ $N$ の文字列 $S_{N}$ が与えられます。 $S$ の $i$ 文字目以降のみからなる部分文字列を $S_{i}$ と表すことにします。 $k = 1, 2, \dots, N$ のそれぞれに対して

$f(S_{k}, S_{1}) + f(S_{k}, S_{2}) + \dots + f(S_{k}, S_{k})$

の値を求めてください。

制約

$1 \le N \le 10^{6}$
$S$ は英小文字のみからなる文字列

Suffix Array

$S$ の suffix (後ろ何文字か) たちについて考える問題自体はよく出題されていて、Suffix Array と呼ばれる配列を作ることで解決できることが多い。

Suffix Array とは、文字列 $S$ の suffix たちを辞書順に並べたもののことである。たとえば $S$ = "abcababaabc" (11 文字) のとき、suffix を集めると

, c, bc, abc, aabc, baabc, abaabc, babaabc, ababaabc, cababaabc, bcababaabc, abcababaabc

の 12 個となる。これらを辞書順に並べると、

aabc
abaabc
ababaabc
abc
abcababaabc
baabc
babaabc
bc
bcababaabc
c
cababaabc

となる。Suffix Array では、次のような配列を管理する。

sa[i] ← $S$ の集合の中で辞書順で小さい順に $i$ 番目のものは $S$ の何文字目から始まるか
rsa[i] ← 「 $S$ の $i$ 文字目以降」が Suffix の集合の中で辞書順で何番目か (sa の逆順列)

これらはともに、SA-IS と呼ばれるアルゴリズムで、 $O(N)$ の計算量で求められる (省略)。上の例では

sa  = {11 7 5 3 8 0 6 4 9 1 10 2}
rsa = {5 9 11 3 7 2 6 1 4 8 10 0}

となる。

高さ配列 LCP

配列 sa を用いていろいろすると、さらに次の高さ配列 lcp が求められる。これが色々と便利なのだ。

lcp[i] ← 文字列 $S$ の suffix を辞書順で小さい順に並べたときの、 $i$ 番目のものと $i+1$ 番目のものについて、先頭が何文字まで一致しているか

たとえば今回の $S$ = abcababaabc の例でいえば、suffix たちを小さい順に並べると

aabc
abaabc
ababaabc
abc
abcababaabc
baabc
babaabc
bc
bcababaabc
c
cababaabc

であり、それぞれの隣接する 2 つの文字列の「先頭が何文字一致しているか」を考えると

lcp = {0 1 3 2 3 0 2 1 2 0 1}

となる。たとえば lcp[2]は、abaabc (辞書順 2 番目) と ababaabc (辞書順 3 番目) の先頭が 3 文字まで一致しているので、lcp[2] = 3 となっている。

さて、上で求めた配列 rsa と、高さ配列 lcp を組み合わせると、次の問に対する見通しがよくなる。

文字列 $S$ において、 $f(S_{i}, S_{j})$ ( $i, j$ は 0-indexed とする) は次のように求められる。

pi = rsa[i] とする ( $S_{i}$ が辞書順で何番目かを pi とする)
pj = rsa[j] とする ( $S_{j}$ が辞書順で何番目かを pj とする)

として、l = min(pi, pj), r = max(pi, pj) とする。このとき $f(S_{i}, S_{j})$ は、配列 lcp において区間 [l, r) の最小値に一致する。

たとえば $i = 5$ (abaabc) と、 $j = 0$ (abcababaabc) について考えてみよう。まず、pi = 2, pj = 5 となる。ここで、辞書順に並べた suffix たちを小さい順に取り出すと

abaabc (2 番目)
ababaabc
abc
abcababaabc (5 番目)

となる。このとき $f(S_{5}, f_{0})$ とは、

2 番目と 3 番目の先頭一致度：3 (= lcp[2])
3 番目と 4 番目の先頭一致度：2 (= lcp[3])
4 番目と 5 番目の先頭一致度：3 (= lcp[4])

の最小値 (最小となるところがボトルネックになる) をとって、2 と求められることがわかる。

もとの問題の言い換え

ここまでくると、元の問題は次のように言い換えられる。

サイズ $N$ の配列 $L_{0}, L_{1}, \dots, L_{N-1}$ (高さ配列 lcp のこと) が与えられる。配列 $L$ の区間 $\lbrack l, r)$ における最小値を $g(l, r)$ と書くことにする。

各 $k = 1, 2, \dots, N$ に対して、

$g(1, k) + g(2, k) + \dots + g(k-1, k) + g(k, k+1) + \dots + g(k, N)$

の値を求めよ。

(なお実際の答えは sk = sa[k] として、上記の答えに N - sk を加算した値を res[sk] に格納することに注意)

なお、rsa はかならず最後尾が 0 になる。なので、もとの文字列での $k = 0, 1, \dots, N-1$ に関する問題は、rsa による変換をかましたあとの配列 L に関する問題では、 $k = 1, 2, \dots, N$ について解いていくことになる。この辺の添字ガチャを合わせるの、苦労した...

言い換え後の問題

このように言い換えたあとも普通に難しい。区間の min をとるのは Sparse Table などを用いて $O(1)$ でできるようにしたとしても、愚直に $g$ の総和を求めると $O(N^{2})$ の計算量となってしまう。

ここではとりあえず、各 $k$ に対して

$h(k) = g(1, k) + g(2, k) + \dots + g(k-1, k)$

を求めていくことにする。反対側の $g(k, k+1) + \dots + g(k, N)$ も同様に求められる。

たとえば $L = (100000, 6, 8, 11, 9, 7, 11, 4, 9)$ として様子を観察してみよう。各 $k$ に対して、ベクトル $(g(1, k), g(2, k), \dots, g(k-1, k))$ は次のようになる。

$k = 1$ のとき、 $()$ なので、 $h(1) = 0$
$k = 2$ のとき、 $(6)$ なので、 $h(2) = 6$
$k = 3$ のとき、 $(6, 8)$ なので、 $h(3) = 14$
$k = 4$ のとき、 $(6, 8, 11)$ なので、 $h(4) = 25$
$k = 5$ のとき、 $(6, 8, 9, 9)$ なので、 $h(5) = 32$
$k = 6$ のとき、 $(6, 7, 7, 7, 7)$ なので、 $h(6) = 34$
$k = 7$ のとき、 $(6, 7, 7, 7, 7, 11)$ なので、 $h(7) = 45$
$k = 8$ のとき、 $(4, 4, 4, 4, 4, 4, 4)$ なので、 $h(8) = 28$
$k = 9$ のとき、 $(4, 4, 4, 4, 4, 4, 4, 9)$ なので、 $h(9) = 37$

つまり、 $L$ の要素が新しく増えるとき、過去へ過去へと「その値よりも小さいところ」にぶつかるまで遡って、その値に書き換えてしまうような操作に対応することがわかる。このような操作は stack を用いてシミュレーションすることができる。

よって、全体をとおして計算量は $O(N)$ となる。

コード

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

// SA-IS (O(N))
template<class Str> struct SuffixArray {
    // data
    Str str;
    vector<int> sa, rsa, lcp;
    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);
        rsa.assign(N + 1, 0);
        for (int i = 0; i <= N; ++i) rsa[sa[i]] = i;
        calc_lcp(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> calc_lcp(const vector<int> &s) {
        int N = (int)s.size();
        lcp.assign(N - 1, 0);
        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;
        }
        return lcp;
    }
};



void solve(int N, const string &S) {
    vector<long long> res(N, 0), left(N + 1, 0), right(N + 1, 0);
    SuffixArray<string> sa(S);
    auto lcp = sa.lcp;

    // (値, index)
    using pll = pair<long long,long long>;
    
    // 左から右へ
    stack<pll> st;
    for (int k = 2; k <= N; ++k) {
        long long val = lcp[k - 1];
        
        // stack から val 以上である限り削除
        while (!st.empty() && st.top().first >= val) st.pop();

        // 答え
        int id = !st.empty() ? st.top().second : 1;
        left[k] = left[id] + val * (k - id);

        // stack に push
        st.push(pll(val, k));
    }

    // 右から左へ
    while (!st.empty()) st.pop();
    for (int k = N - 1; k >= 1; --k) {
        long long val = lcp[k];

        // stack から val 以上である限り削除
        while (!st.empty() && st.top().first >= val) st.pop();

        // 答え
        int id = !st.empty() ? st.top().second : N;
        right[k] = right[id] + val * (id - k);

        // stack に push
        st.push(pll(val, k));
    }

    // 答え
    for (int k = 1; k <= N; ++k) {
        res[sa[k]] = left[k] + right[k] + (N - sa[k]);
    }
    for (auto v: res) cout << v << endl;
}

int main() {
    int N;
    string S;
    while (cin >> N >> S) solve(N, S);
}