Suffix Array Construction

suffix_array.cc

struct Suffix {
    int idx, rank0, rank1;
};

bool suffix_lt(Suffix& a, Suffix& b)
{
    if (a.rank0 == b.rank0) return a.rank1 < b.rank1;
    return a.rank0 < b.rank0;
}

V<int> suffix_array(string& s)
{
    int n = s.size();
    int m = log(n)/log(2);
    V<V<int>> p(m+1, V<int>(n));
    V<Suffix> l(n);
    for (int i = 0; i < n; ++i) {
        p[0][i] = s[i] - 'A';
    }
    for (int k = 1, cnt = 1; k <= m; ++k, cnt *= 2) {
        for (int i = 0; i < n; ++i) {
            l[i].idx = i;
            l[i].rank0 = p[k-1][i];
            l[i].rank1 = (i+cnt < n ? p[k-1][i+cnt] : -1);
        }
        sort(l.begin(), l.end(), suffix_lt);
        for (int i = 0; i < n; ++i) {
            int eq_to_last = i > 0 and l[i].rank0 == l[i-1].rank0 and l[i].rank1 == l[i-1].rank1;
            p[k][l[i].idx] = (eq_to_last ? p[k][l[i-1].idx] : i);
        }
    }
    V<int> sa(n);
    for (int i = 0; i < n; ++i) {
        sa[i] = l[i].idx;
    }
    return sa;
}