String/SuffixArray.hpp

• View this file on GitHub
• Last update: 2023-02-11 22:28:05+08:00
• Include: #include "String/SuffixArray.hpp"

Depends on

• DataStructure/RMQ.hpp
• Template/Template.hpp

Verified with

• Verify/SuffixArray.test.cpp

Code

#pragma once
#include "../DataStructure/RMQ.hpp"
#include "../Template/Template.hpp"

struct SuffixArray {
    // lc[i] ->lcp(sa[i],sa[i+1])
    int n;
    std::vector<int> sa, rk, lc;
    RMQ<int> *rmq;
    SuffixArray(const std::string &s) : rmq(nullptr) {
        n = s.length();
        sa.resize(n);
        lc.resize(n - 1);
        rk.resize(n);
        std::iota(sa.begin(), sa.end(), 0);
        std::sort(sa.begin(), sa.end(),
                  [&](int a, int b) { return s[a] < s[b]; });
        rk[sa[0]] = 0;
        for (int i = 1; i < n; ++i)
            rk[sa[i]] = rk[sa[i - 1]] + (s[sa[i]] != s[sa[i - 1]]);
        int k = 1;
        std::vector<int> tmp, cnt(n);
        tmp.reserve(n);
        while (rk[sa[n - 1]] < n - 1) {
            tmp.clear();
            for (int i = 0; i < k; ++i) tmp.push_back(n - k + i);
            for (auto i : sa)
                if (i >= k) tmp.push_back(i - k);
            std::fill(cnt.begin(), cnt.end(), 0);
            for (int i = 0; i < n; ++i) ++cnt[rk[i]];
            for (int i = 1; i < n; ++i) cnt[i] += cnt[i - 1];
            for (int i = n - 1; i >= 0; --i) sa[--cnt[rk[tmp[i]]]] = tmp[i];
            std::swap(rk, tmp);
            rk[sa[0]] = 0;
            for (int i = 1; i < n; ++i)
                rk[sa[i]] =
                    rk[sa[i - 1]] + (tmp[sa[i - 1]] < tmp[sa[i]] ||
                                     sa[i - 1] + k == n ||
                                     tmp[sa[i - 1] + k] < tmp[sa[i] + k]);
            k *= 2;
        }
        for (int i = 0, j = 0; i < n; ++i) {
            if (rk[i] == 0) {
                j = 0;
            } else {
                for (j -= j > 0; i + j < n && sa[rk[i] - 1] + j < n &&
                                 s[i + j] == s[sa[rk[i] - 1] + j];)
                    ++j;
                lc[rk[i] - 1] = j;
            }
        }
        if (n > 1) {
            rmq = new RMQ(lc);
        }
    }
    ~SuffixArray() {
        if (rmq) {
            delete rmq;
        }
    }
    int lcp(int x, int y) {
        if (x == n || y == n) {
            return 0;
        }
        if (x == y) {
            return n - x;
        }
        x = rk[x];
        y = rk[y];
        if (x > y) {
            std::swap(x, y);
        }
        return rmq->rangeMin(x, y);
    }
};

#line 2 "Template/Template.hpp"

#include <bits/stdc++.h>

using i64 = std::int64_t;
#line 3 "DataStructure/RMQ.hpp"

template <class T, class Cmp = std::less<T>>
struct RMQ {
    const int n;
    const Cmp cmp;
    std::vector<std::vector<T>> a;
    RMQ(const std::vector<T> &init) : n(init.size()), cmp(Cmp()) {
        int lg = std::__lg(n);
        a.assign(lg + 1, std::vector<T>(n));
        for (int j = 0; j <= lg; j++) {
            for (int i = 0; i + (1 << j) <= n; i++) {
                a[j][i] =
                    (j == 0 ? init[i]
                            : std::min(a[j - 1][i],
                                       a[j - 1][i + (1 << (j - 1))], cmp));
            }
        }
    }
    T rangeMin(int l, int r) {
        int k = std::__lg(r - l);
        return std::min(a[k][l], a[k][r - (1 << k)], cmp);
    }
};
#line 4 "String/SuffixArray.hpp"

struct SuffixArray {
    // lc[i] ->lcp(sa[i],sa[i+1])
    int n;
    std::vector<int> sa, rk, lc;
    RMQ<int> *rmq;
    SuffixArray(const std::string &s) : rmq(nullptr) {
        n = s.length();
        sa.resize(n);
        lc.resize(n - 1);
        rk.resize(n);
        std::iota(sa.begin(), sa.end(), 0);
        std::sort(sa.begin(), sa.end(),
                  [&](int a, int b) { return s[a] < s[b]; });
        rk[sa[0]] = 0;
        for (int i = 1; i < n; ++i)
            rk[sa[i]] = rk[sa[i - 1]] + (s[sa[i]] != s[sa[i - 1]]);
        int k = 1;
        std::vector<int> tmp, cnt(n);
        tmp.reserve(n);
        while (rk[sa[n - 1]] < n - 1) {
            tmp.clear();
            for (int i = 0; i < k; ++i) tmp.push_back(n - k + i);
            for (auto i : sa)
                if (i >= k) tmp.push_back(i - k);
            std::fill(cnt.begin(), cnt.end(), 0);
            for (int i = 0; i < n; ++i) ++cnt[rk[i]];
            for (int i = 1; i < n; ++i) cnt[i] += cnt[i - 1];
            for (int i = n - 1; i >= 0; --i) sa[--cnt[rk[tmp[i]]]] = tmp[i];
            std::swap(rk, tmp);
            rk[sa[0]] = 0;
            for (int i = 1; i < n; ++i)
                rk[sa[i]] =
                    rk[sa[i - 1]] + (tmp[sa[i - 1]] < tmp[sa[i]] ||
                                     sa[i - 1] + k == n ||
                                     tmp[sa[i - 1] + k] < tmp[sa[i] + k]);
            k *= 2;
        }
        for (int i = 0, j = 0; i < n; ++i) {
            if (rk[i] == 0) {
                j = 0;
            } else {
                for (j -= j > 0; i + j < n && sa[rk[i] - 1] + j < n &&
                                 s[i + j] == s[sa[rk[i] - 1] + j];)
                    ++j;
                lc[rk[i] - 1] = j;
            }
        }
        if (n > 1) {
            rmq = new RMQ(lc);
        }
    }
    ~SuffixArray() {
        if (rmq) {
            delete rmq;
        }
    }
    int lcp(int x, int y) {
        if (x == n || y == n) {
            return 0;
        }
        if (x == y) {
            return n - x;
        }
        x = rk[x];
        y = rk[y];
        if (x > y) {
            std::swap(x, y);
        }
        return rmq->rangeMin(x, y);
    }
};

Back to top page