めっちゃ面白い問題だった！　Suffix Array の練習問題。

問題へのリンク

問題概要

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

$S$ の suffix は空文字列を除いて $N$ 個ある。これらの suffix を適切な順序で連結させて 1 つの文字列を作るとき、それが辞書順最小となるような順序を求めよ。

制約

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

考えたこと

まずこの問題の前提となる知見として、文字列 $a_{0}, a_{1}, \dots, a_{N-1}$ を適切な順序で連結して辞書順最小になるようにする方法は、連結させたどの隣接する文字列 $p, q$ に対しても

$p + q \lt q + p$

が成り立つようにすることだ。このことの証明は次の記事にて。

drken1215.hatenablog.com

ここで、 $S$ の $i$ 文字目以降の文字列 (S.substr(i) のこと) を $S_{i}$ と書くことにする。 $S_{i}$ の長さは $|S_{i}| = N - i$ となる。

一般に 2 つの suffix $S_{i}, S_{j}$ に対して

$S_{i} + S_{j} \lt S_{j} + S_{i}$

となる条件を考察しよう。簡単のため、 $i \lt j$ と仮定する。 $j \lt i$ の場合は比較関数をひっくり返せばよい。 $i \lt j$ のとき、 $|S_{i}| \gt |S_{j}|$ となることに注意しておく。

$S_{j}$ が $S_{i}$ の prefix になっていないとき

$S_{j}$ が $S_{i}$ の prefix になっているかどうかの判定は、lcp(i, j) == N-j かどうかによって判定できる。まず、 $S_{j}$ が $S_{i}$ の prefix になっていない場合を考察する。

この場合は下図のように、 $S_{i} + S_{j}$ と $S_{j} + S_{i}$ とを比べたときに、先頭から長さ $|S_{j}|$ 以内の範囲で、文字が異なる箇所がある。

よって、 $S_{i} + S_{j}$ と $S_{j} + S_{i}$ の辞書順の大小関係は、単純に $S_{i}$ と $S_{j}$ の辞書順の大小関係に一致する。つまり、Suffix Array のランク配列を rank として、

rank[i] < rank[j]

によって判定できる。

$S_{j}$ が $S_{i}$ の prefix になっているとき

この場合が要注意だ。

この場合は $l = |S_{j}| (= N-j)$ として、下図で青く示したように、 $S_{i} + S_{j}$ と $S_{j} + S_{i}$ の先頭の $l$ 文字と末尾の $l$ 文字は一致することになる。末尾が一致することについては suffix の定義から従う。

よって、先頭と末尾の $l$ 文字分を除外して残された部分は

$S_{i} + S_{j}$ → $S_{i+l}$
$S_{i} + S_{j}$ → $S_{i}$ の先頭から $l$ 文字分

となる。よって

$S_{i+l} \gt S_{i}$ である場合は、確実に $S_{i} + S_{j} \gt S_{j} + S_{i}$ であると言える
である場合は、
- 「 $S_{i+l} \lt S_{i}$ の先頭から $l$ 文字分」である場合は $S_{i} + S_{j} \lt S_{j} + S_{i}$ である
- 「 $S_{i+l} = S_{i}$ の先頭から $l$ 文字分」である場合は、 $S_{i} + S_{j} = S_{j} + S_{i}$ である

「 $=$ 」になる場合が微妙ではあるが、まとめると

$S_{i+l} \lt S_{i}$ 　 $\Leftrightarrow$ 　 $S_{i} + S_{j} \le S_{j} + S_{i}$

であると言える。よって、比較関数は

rank[i+l] < rank[i]

と実装すれば問題ない。

コード

上述の比較関数に基づいて、 $0, 1, \dots, N-1$ をソートすればよい。計算量は

Suffix Array を構築して rank を求める： $O(N)$
比較関数の処理：lcp を求める部分に Sparse Table を用いると、前処理 $O(N \log N)$ 、各比較処理は $O(1)$ でできる
ソートの実行： $O(N \log N)$ でできる

以上より、全体として $O(N \log N)$ の計算量で求められる。

#include <bits/stdc++.h>
using namespace std;
using pint = pair<int,int>;

// Sparse Table
template<class MeetSemiLattice> struct SparseTable {
    vector<vector<MeetSemiLattice> > dat;
    vector<int> height;
    
    SparseTable() { }
    SparseTable(const vector<MeetSemiLattice> &vec) { init(vec); }
    void init(const vector<MeetSemiLattice> &vec) {
        int n = (int)vec.size(), h = 0;
        while ((1<<h) < n) ++h;
        dat.assign(h, vector<MeetSemiLattice>(1<<h));
        height.assign(n+1, 0);
        for (int i = 2; i <= n; i++) height[i] = height[i>>1]+1;
        for (int i = 0; i < n; ++i) dat[0][i] = vec[i];
        for (int i = 1; i < h; ++i)
            for (int j = 0; j < n; ++j)
                dat[i][j] = min(dat[i-1][j], dat[i-1][min(j+(1<<(i-1)),n-1)]);
    }
    
    MeetSemiLattice get(int a, int b) {
        return min(dat[height[b-a]][a], dat[height[b-a]][b-(1<<height[b-a])]);
    }
};

// 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> rank;  // rank[sa[i]] = i
    vector<int> lcp;   // lcp[i]: the lcp of sa[i] and sa[i+1] (i = 0, 1, ..., n-1)
    SparseTable<int> st;  // use for calcultating lcp(i, j)

    // getter
    int& operator [] (int i) {
        return sa[i];
    }
    vector<int> get_sa() { return sa; }
    vector<int> get_rank() { return rank; }
    vector<int> get_lcp() { return lcp; }

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

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

    // find min id that str.substr(sa[id]) >= T
    int lower_bound(const Str& T) {
        int left = -1, right = sa.size();
        while (right - left > 1) {
            int mid = (left + right) / 2;
            if (str.compare(sa[mid], string::npos, T) < 0)
                left = mid;
            else
                right = mid;
        }
        return right;
    }

    // find min id that str.substr(sa[id], T.size()) > T
    int upper_bound(const Str& T) {
        int left = -1, right = sa.size();
        while (right - left > 1) {
            int mid = (left + right) / 2;
            if (str.compare(sa[mid], T.size(), T) <= 0)
                left = mid;
            else
                right = mid;
        }
        return right;
    }

    // search
    bool is_contain(const Str& T) {
        int lb = lower_bound(T);
        if (lb >= sa.size()) return false;
        return str.compare(sa[lb], T.size(), T) == 0;
    }

    // find lcp
    void calcLCP(const vector<int> &s) {
        int N = (int)s.size();
        rank.assign(N, 0), lcp.assign(N, 0);
        for (int i = 0; i < N; ++i) rank[sa[i]] = i;
        int h = 0;
        for (int i = 0; i < N - 1; ++i) {
            int pi = sa[rank[i] - 1];
            if (h > 0) --h;
            for (; pi + h < N && i + h < N; ++h) {
                if (s[pi + h] != s[i + h]) break;
            }
            lcp[rank[i] - 1] = h;
        }
    }
    
    // build sparse table for calculating lcp
    void buildSparseTable() {
        st.init(lcp);
    }

    // calc lcp of str.sutstr(a) and str.substr(b)
    int getLCP(int a, int b) {
        return st.get(min(rank[a], rank[b]), max(rank[a], rank[b]));
    }
};

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

    // Suffix Array の構築
    SuffixArray<string> suf(S);
    vector<int> rank = suf.get_rank();

    // suffix の比較
    function<bool (int, int)> cmp = [&](int i, int j) {
        // i < j を仮定できるようにする
        if (i == j) return true;
        if (i > j) return not cmp(j, i);

        // i と j の lcp
        int len = suf.getLCP(i, j);
        
        if (N-j > len) {
            // i と j が prefix の関係にないならば小さい順
            return rank[i] < rank[j];
        } else {
            // i の len 文字目以降と i を比較
            return rank[i+len] < rank[i];
        }
    };

    // ソート
    vector<int> ids(N);
    for (int i = 0; i < N; ++i) ids[i] = i;
    sort(ids.begin(), ids.end(), cmp);
    for (auto v : ids) cout << v+1 << endl;
}