cp-book
ecnerwala's competitive programming library
#include "suffix_array.hpp"
Depends on
Verified with
Code
| Coverage | Exec / Excl / Total | |
|---|---|---|
| Lines | 75.3% | 183 / 0 / 243 |
| Functions | 81.3% | 26 / 0 / 32 |
| Branches | 93.5% | 376 / 0 / 402 |
| Full report |
#pragma once
/*
* This is mostly inspired by https://golang.org/src/index/suffixarray/sais.go.
*/
#include <algorithm>
#include <vector>
#include <string>
#include <cassert>
#include <cstring>
#include <type_traits>
#include "rmq.hpp"
template<class T> int sz(T&& arg) { using std::size; return int(size(std::forward<T>(arg))); }
// Layered suffix array: SuffixArrayBase computes just sa/rank, each further
// layer statically opts into one more derived structure. Use the leaf classes
// SuffixArray, SuffixArrayLCP, or SuffixArrayRMQ; the named constructors on
// each return that type.
template <typename Self> class SuffixArrayBase {
public:
using index_t = int;
int N;
std::vector<index_t> sa;
std::vector<index_t> rank;
SuffixArrayBase() : N(0) {}
template <typename String> static Self construct_raw(const String& S, index_t sigma) {
Self res;
res.build(S, sigma);
return res;
}
// Pass a function which returns a value in [0, sigma)
template <typename String, typename F> static Self map_and_construct(const String& S, const F& f, int sigma) {
std::vector<decltype((f(S[0])))> mapped(sz(S));
for (int i = 0; i < sz(S); i++) {
mapped[i] = f(S[i]);
assert(0 <= int(mapped[i]) && int(mapped[i]) < sigma);
}
return construct_raw(mapped, sigma);
}
// Sorts the elements of S and then runs suffix array. This takes O(N log N) time with no dependence on sigma.
template <typename String> static Self sort_and_construct(const String& S) {
using std::begin;
using std::end;
using value_type = typename std::iterator_traits<decltype(begin(S))>::value_type;
using compressed_value_type = typename std::conditional<
sizeof(value_type) < sizeof(index_t),
value_type,
index_t
>::type;
std::vector<compressed_value_type> compressed_s(sz(S));
int sigma = 0;
{
std::vector<value_type> vals(begin(S), end(S));
std::sort(vals.begin(), vals.end());
vals.resize(unique(vals.begin(), vals.end()) - vals.begin());
for (int i = 0; i < sz(S); i++) {
compressed_s[i] = compressed_value_type(index_t(std::lower_bound(vals.begin(), vals.end(), S[i]) - vals.begin()));
}
sigma = int(vals.size());
}
return construct_raw(compressed_s, sigma);
}
// Shifts the elements so that sigma = max(S) - min(S) + 1
template <typename String> static Self shift_and_construct(const String& S) {
using std::begin;
using std::end;
using value_type = typename std::iterator_traits<decltype(begin(S))>::value_type;
std::vector<value_type> compressed_s(sz(S));
int sigma = 0;
if (sz(S) > 0) {
value_type lo = *begin(S), hi = *begin(S);
for (const auto& x : S) {
if (x < lo) lo = x;
if (x > hi) hi = x;
}
for (int i = 0; i < sz(S); i++) {
compressed_s[i] = value_type(S[i] - lo);
}
sigma = int(hi - lo + 1);
}
return construct_raw(compressed_s, sigma);
}
// Renumber/filter to only the used elements with bucket sorting. Still takes O(max(S) - min(S) + 1) memory/time,
// but should be less memory than `shift_and_construct` when sigma ~ N and max(S) - min(S) + 1 > N.
template <typename String> static Self bucket_and_construct(const String& S) {
using std::begin;
using std::end;
using value_type = typename std::iterator_traits<decltype(begin(S))>::value_type;
using compressed_value_type = typename std::conditional<
sizeof(value_type) < sizeof(index_t),
value_type,
index_t
>::type;
std::vector<compressed_value_type> compressed_s(sz(S));
int sigma = 0;
if (sz(S) > 0) {
value_type lo = *begin(S), hi = *begin(S);
for (const auto& x : S) {
if (x < lo) lo = x;
if (x > hi) hi = x;
}
std::vector<compressed_value_type> buckets(hi - lo + 1, 0);
for (const auto& x : S) {
buckets[x - lo] = 1;
}
for (int v = 0; v < int(buckets.size()); v++) {
if (buckets[v]) buckets[v] = compressed_value_type(sigma++);
}
for (int i = 0; i < sz(S); i++) {
compressed_s[i] = buckets[S[i] - lo];
}
}
return construct_raw(compressed_s, sigma);
}
protected:
template <typename String> void build(const String& S, index_t sigma) {
N = sz(S);
build_sa(S, sigma);
build_rank();
}
private:
template <typename String> void build_sa(const String& S, index_t sigma) {
sa = std::vector<index_t>(N+1);
assert(sigma >= 0);
for (auto s : S) assert(0 <= index_t(s) && index_t(s) < sigma);
std::vector<index_t> tmp(sigma + std::max(N, sigma));
SuffixArrayBase::sais<String>(N, S, sa.data(), sigma, tmp.data());
}
template <typename String> static void sais(int N, const String& S, index_t* sa, int sigma, index_t* tmp) {
if (N == 0) {
sa[0] = 0;
return;
} else if (N == 1) {
sa[0] = 1;
sa[1] = 0;
return;
}
// Phase 1: Initialize the frequency array, which will let us lookup buckets.
index_t* freq = tmp; tmp += sigma;
memset(freq, 0, sizeof(*freq) * sigma);
for (int i = 0; i < N; i++) {
++freq[index_t(S[i])];
}
auto build_bucket_start = [&]() {
int cur = 1;
for (int v = 0; v < sigma; v++) {
tmp[v] = cur;
cur += freq[v];
}
};
auto build_bucket_end = [&]() {
int cur = 1;
for (int v = 0; v < sigma; v++) {
cur += freq[v];
tmp[v] = cur;
}
};
int num_pieces = 0;
int first_endpoint = 0;
// Phase 2: find the right-endpoints of the pieces
{
build_bucket_end();
// Initialize the final endpoint out-of-band this way so that we don't try to look up tmp[-1].
// This doesn't count towards num_pieces.
sa[0] = N;
index_t c0 = S[N-1], c1 = -1; bool isS = false;
for (int i = N-2; i >= 0; i--) {
c1 = c0;
c0 = S[i];
if (c0 < c1) {
isS = true;
} else if (c0 > c1 && isS) {
isS = false;
// insert i+1
sa[first_endpoint = --tmp[c1]] = i+1;
++num_pieces;
}
}
}
// If num_pieces <= 1, we don't need to actually run the recursion, it's just sorted automatically
// Otherwise, we're going to rebucket
if (num_pieces > 1) {
// Remove the first endpoint, we don't need to run the IS on this
sa[first_endpoint] = 0;
// Run IS for L-type
{
build_bucket_start();
for (int z = 0; z <= N; z++) {
int v = sa[z];
if (!v) continue;
// Leave for the S-round
if (v < 0) continue;
// clear out our garbage
sa[z] = 0;
--v;
index_t c0 = S[v-1], c1 = S[v];
sa[tmp[c1]++] = (c0 < c1) ? ~v : v;
}
}
index_t* const sa_end = sa + N + 1;
index_t* pieces = sa_end;
// Run IS for S-type and compactify
{
build_bucket_end();
for (int z = N; z >= 0; z--) {
int v = sa[z];
if (!v) continue;
// clear our garbage
sa[z] = 0;
if (v > 0) {
*--pieces = v;
continue;
}
v = ~v;
--v;
index_t c0 = S[v-1], c1 = S[v];
sa[--tmp[c1]] = (c0 > c1) ? v : ~v;
}
}
// Compute the lengths of the pieces in preparation for equality
// comparison, and store them in sa[v/2]. We set the length of the
// final piece to 0; it compares unequal to everything because of
// the sentinel.
{
int prv_start = N;
index_t c0 = S[N-1], c1 = -1; bool isS = false;
for (int i = N-2; i >= 0; i--) {
c1 = c0;
c0 = S[i];
if (c0 < c1) {
isS = true;
} else if (c0 > c1 && isS) {
isS = false;
// insert i+1
int v = i+1;
sa[v>>1] = prv_start == N ? 0 : prv_start - v;
prv_start = v;
}
}
}
// Compute the alphabet, storing the result into sa[v/2].
int next_sigma = 0;
{
int prv_len = -1, prv_v = 0;
for (int i = 0; i < num_pieces; i++) {
int v = pieces[i];
int len = sa[v>>1];
bool eq = prv_len == len;
for (int a = 0; eq && a < len; ++a) {
eq = S[v+a] == S[prv_v+a];
}
if (!eq) {
next_sigma++;
prv_len = len;
prv_v = v;
}
sa[v>>1] = next_sigma; // purposely leave this 1 large to check != 0
}
}
if (next_sigma == num_pieces) {
sa[0] = N;
memcpy(sa+1, pieces, sizeof(*sa) * num_pieces);
} else {
index_t* next_S = sa_end;
// Finally, pack the input to the SA
{
for (int i = (N-1)>>1; i >= 0; i--) {
int v = sa[i];
if (v) *--next_S = v-1;
sa[i] = 0;
}
}
memset(sa, 0, sizeof(*sa) * (num_pieces+1));
sais<const index_t*>(num_pieces, next_S, sa, next_sigma, tmp);
{ // Compute the piece start points again and use those to map up the suffix array
next_S = sa_end;
index_t c0 = S[N-1], c1 = -1; bool isS = false;
for (int i = N-2; i >= 0; i--) {
c1 = c0;
c0 = S[i];
if (c0 < c1) {
isS = true;
} else if (c0 > c1 && isS) {
isS = false;
int v = i+1;
*--next_S = v;
}
}
sa[0] = N;
for (int i = 1; i <= num_pieces; i++) {
sa[i] = next_S[sa[i]];
}
}
}
// zero everything else
memset(sa+num_pieces+1, 0, sizeof(*sa) * (N - num_pieces));
{
// Scatter the finished pieces
build_bucket_end();
for (int i = num_pieces; i > 0; i--) {
int v = sa[i];
sa[i] = 0;
index_t c1 = S[v];
sa[--tmp[c1]] = v;
}
}
}
// Home stretch! Just finish out with the L-type and then S-type
{
build_bucket_start();
for (int z = 0; z <= N; z++) {
int v = sa[z];
if (v <= 0) continue;
--v;
index_t c1 = S[v];
index_t c0 = v ? S[v-1] : c1; // if v = 0, we don't want to invert
sa[tmp[c1]++] = (c0 < c1) ? ~v : v;
}
}
// This just aggressively overwrites our original scattered pieces with the correct values
{
build_bucket_end();
for (int z = N; z >= 0; z--) {
int v = sa[z];
if (v >= 0) continue;
sa[z] = v = ~v;
--v;
index_t c1 = S[v];
index_t c0 = v ? S[v-1] : c1+1;
sa[--tmp[c1]] = (c0 > c1) ? v : ~v;
}
}
}
void build_rank() {
rank = std::vector<index_t>(N+1);
for (int i = 0; i <= N; i++) rank[sa[i]] = i;
}
};
class SuffixArray : public SuffixArrayBase<SuffixArray> {};
template <typename Self> class SuffixArrayLCPBase : public SuffixArrayBase<Self> {
public:
using index_t = typename SuffixArrayBase<Self>::index_t;
// lcp[i] = lcp(sa[i], sa[i+1])
std::vector<index_t> lcp;
protected:
friend SuffixArrayBase<Self>;
template <typename String> void build(const String& S, index_t sigma) {
SuffixArrayBase<Self>::build(S, sigma);
build_lcp(S);
}
private:
template <typename String> void build_lcp(const String& S) {
int N = this->N;
const auto& sa = this->sa;
const auto& rank = this->rank;
assert(sz(S) == N);
lcp = std::vector<index_t>(N);
for (int i = 0, k = 0; i < N - 1; i++) {
int j = sa[rank[i]-1];
while (k < N - std::max(i, j) && S[i+k] == S[j+k]) k++;
lcp[rank[i]-1] = k;
if (k) --k;
}
}
};
class SuffixArrayLCP : public SuffixArrayLCPBase<SuffixArrayLCP> {};
template <typename Self> class SuffixArrayRMQBase : public SuffixArrayLCPBase<Self> {
public:
using index_t = typename SuffixArrayLCPBase<Self>::index_t;
RangeMinQuery<std::pair<index_t, index_t>> rmq;
index_t get_lcp(index_t a, index_t b) const {
if (a == b) return this->N-a;
a = this->rank[a], b = this->rank[b];
if (a > b) std::swap(a, b);
return rmq.query(a, b-1).first;
}
// Get the split in the suffix tree, using half-open intervals
// Returns len, idx
std::pair<index_t, index_t> get_split(index_t l, index_t r) const {
assert(r - l > 1);
return rmq.query(l, r-2);
}
protected:
friend SuffixArrayBase<Self>;
template <typename String> void build(const String& S, index_t sigma) {
SuffixArrayLCPBase<Self>::build(S, sigma);
build_rmq();
}
private:
void build_rmq() {
int N = this->N;
const auto& lcp = this->lcp;
std::vector<std::pair<index_t, index_t>> lcp_idx(N);
for (int i = 0; i < N; i++) {
lcp_idx[i] = {lcp[i], i+1};
}
rmq = RangeMinQuery<std::pair<index_t, index_t>>(std::move(lcp_idx));
}
};
class SuffixArrayRMQ : public SuffixArrayRMQBase<SuffixArrayRMQ> {};
class PrefixArrayRMQ : private SuffixArrayRMQ {
PrefixArrayRMQ(const SuffixArrayRMQ& sa_) : SuffixArrayRMQ(sa_) {}
PrefixArrayRMQ(SuffixArrayRMQ&& sa_) : SuffixArrayRMQ(std::move(sa_)) {}
public:
PrefixArrayRMQ() {}
template <typename String> static PrefixArrayRMQ construct_raw(const String& S, int sigma) {
return PrefixArrayRMQ(SuffixArrayRMQ::construct_raw(String(S.rbegin(), S.rend()), sigma));
}
// TODO: Fill in other constructors
int get_lcs(int a, int b) const {
return SuffixArrayRMQ::get_lcp(N - a, N - b);
}
};
#include <algorithm>
#include <vector>
#include <string>
#include <cassert>
#include <cstring>
#include <type_traits>
#include <functional>
#include <cstdint>
#line 2 "src/suffix_array.hpp"
/*
* This is mostly inspired by https://golang.org/src/index/suffixarray/sais.go.
*/
#line 13 "src/suffix_array.hpp"
#line 2 "src/rmq.hpp"
#line 7 "src/rmq.hpp"
template <typename T, class Compare = std::less<T>> class RangeMinQuery : private Compare {
static const int BUCKET_SIZE = 32;
static const int BUCKET_SIZE_LOG = 5;
static_assert(BUCKET_SIZE == (1 << BUCKET_SIZE_LOG), "BUCKET_SIZE should be a power of 2");
static const int CACHE_LINE_ALIGNMENT = 64;
int n = 0;
std::vector<T> data;
std::vector<T> pref_data;
std::vector<T> suff_data;
std::vector<T> sparse_table;
std::vector<uint32_t> range_mask;
private:
int num_buckets() const {
return n >> BUCKET_SIZE_LOG;
}
int num_levels() const {
return num_buckets() ? 32 - __builtin_clz(num_buckets()) : 0;
}
int sparse_table_size() const {
return num_buckets() * num_levels();
}
private:
const T& min(const T& a, const T& b) const {
return Compare::operator()(a, b) ? a : b;
}
void setmin(T& a, const T& b) const {
if (Compare::operator()(b, a)) a = b;
}
template <typename Vec> static int get_size(const Vec& v) { using std::size; return int(size(v)); }
public:
RangeMinQuery() {}
template <typename Vec> explicit RangeMinQuery(const Vec& data_, const Compare& comp_ = Compare())
: Compare(comp_)
, n(get_size(data_))
, data(n)
, pref_data(n)
, suff_data(n)
, sparse_table(sparse_table_size())
, range_mask(n)
{
for (int i = 0; i < n; i++) data[i] = data_[i];
for (int i = 0; i < n; i++) {
if (i & (BUCKET_SIZE-1)) {
uint32_t m = range_mask[i-1];
while (m && !Compare::operator()(data[(i | (BUCKET_SIZE-1)) - __builtin_clz(m)], data[i])) {
m -= uint32_t(1) << (BUCKET_SIZE - 1 - __builtin_clz(m));
}
m |= uint32_t(1) << (i & (BUCKET_SIZE - 1));
range_mask[i] = m;
} else {
range_mask[i] = 1;
}
}
for (int i = 0; i < n; i++) {
pref_data[i] = data[i];
if (i & (BUCKET_SIZE-1)) {
setmin(pref_data[i], pref_data[i-1]);
}
}
for (int i = n-1; i >= 0; i--) {
suff_data[i] = data[i];
if (i+1 < n && ((i+1) & (BUCKET_SIZE-1))) {
setmin(suff_data[i], suff_data[i+1]);
}
}
for (int i = 0; i < num_buckets(); i++) {
sparse_table[i] = data[i * BUCKET_SIZE];
for (int v = 1; v < BUCKET_SIZE; v++) {
setmin(sparse_table[i], data[i * BUCKET_SIZE + v]);
}
}
for (int l = 0; l+1 < num_levels(); l++) {
for (int i = 0; i + (1 << (l+1)) <= num_buckets(); i++) {
sparse_table[(l+1) * num_buckets() + i] = min(sparse_table[l * num_buckets() + i], sparse_table[l * num_buckets() + i + (1 << l)]);
}
}
}
T query(int l, int r) const {
assert(l <= r);
int bucket_l = (l >> BUCKET_SIZE_LOG);
int bucket_r = (r >> BUCKET_SIZE_LOG);
if (bucket_l == bucket_r) {
uint32_t msk = range_mask[r] & ~((uint32_t(1) << (l & (BUCKET_SIZE-1))) - 1);
int ind = (l & ~(BUCKET_SIZE-1)) + __builtin_ctz(msk);
return data[ind];
} else {
T ans = min(suff_data[l], pref_data[r]);
bucket_l++;
if (bucket_l < bucket_r) {
int level = (32 - __builtin_clz(bucket_r - bucket_l)) - 1;
setmin(ans, sparse_table[level * num_buckets() + bucket_l]);
setmin(ans, sparse_table[level * num_buckets() + bucket_r - (1 << level)]);
}
return ans;
}
}
};
template <typename T> using RangeMaxQuery = RangeMinQuery<T, std::greater<T>>;
#line 15 "src/suffix_array.hpp"
template<class T> int sz(T&& arg) { using std::size; return int(size(std::forward<T>(arg))); }
// Layered suffix array: SuffixArrayBase computes just sa/rank, each further
// layer statically opts into one more derived structure. Use the leaf classes
// SuffixArray, SuffixArrayLCP, or SuffixArrayRMQ; the named constructors on
// each return that type.
template <typename Self> class SuffixArrayBase {
public:
using index_t = int;
int N;
std::vector<index_t> sa;
std::vector<index_t> rank;
SuffixArrayBase() : N(0) {}
template <typename String> static Self construct_raw(const String& S, index_t sigma) {
Self res;
res.build(S, sigma);
return res;
}
// Pass a function which returns a value in [0, sigma)
template <typename String, typename F> static Self map_and_construct(const String& S, const F& f, int sigma) {
std::vector<decltype((f(S[0])))> mapped(sz(S));
for (int i = 0; i < sz(S); i++) {
mapped[i] = f(S[i]);
assert(0 <= int(mapped[i]) && int(mapped[i]) < sigma);
}
return construct_raw(mapped, sigma);
}
// Sorts the elements of S and then runs suffix array. This takes O(N log N) time with no dependence on sigma.
template <typename String> static Self sort_and_construct(const String& S) {
using std::begin;
using std::end;
using value_type = typename std::iterator_traits<decltype(begin(S))>::value_type;
using compressed_value_type = typename std::conditional<
sizeof(value_type) < sizeof(index_t),
value_type,
index_t
>::type;
std::vector<compressed_value_type> compressed_s(sz(S));
int sigma = 0;
{
std::vector<value_type> vals(begin(S), end(S));
std::sort(vals.begin(), vals.end());
vals.resize(unique(vals.begin(), vals.end()) - vals.begin());
for (int i = 0; i < sz(S); i++) {
compressed_s[i] = compressed_value_type(index_t(std::lower_bound(vals.begin(), vals.end(), S[i]) - vals.begin()));
}
sigma = int(vals.size());
}
return construct_raw(compressed_s, sigma);
}
// Shifts the elements so that sigma = max(S) - min(S) + 1
template <typename String> static Self shift_and_construct(const String& S) {
using std::begin;
using std::end;
using value_type = typename std::iterator_traits<decltype(begin(S))>::value_type;
std::vector<value_type> compressed_s(sz(S));
int sigma = 0;
if (sz(S) > 0) {
value_type lo = *begin(S), hi = *begin(S);
for (const auto& x : S) {
if (x < lo) lo = x;
if (x > hi) hi = x;
}
for (int i = 0; i < sz(S); i++) {
compressed_s[i] = value_type(S[i] - lo);
}
sigma = int(hi - lo + 1);
}
return construct_raw(compressed_s, sigma);
}
// Renumber/filter to only the used elements with bucket sorting. Still takes O(max(S) - min(S) + 1) memory/time,
// but should be less memory than `shift_and_construct` when sigma ~ N and max(S) - min(S) + 1 > N.
template <typename String> static Self bucket_and_construct(const String& S) {
using std::begin;
using std::end;
using value_type = typename std::iterator_traits<decltype(begin(S))>::value_type;
using compressed_value_type = typename std::conditional<
sizeof(value_type) < sizeof(index_t),
value_type,
index_t
>::type;
std::vector<compressed_value_type> compressed_s(sz(S));
int sigma = 0;
if (sz(S) > 0) {
value_type lo = *begin(S), hi = *begin(S);
for (const auto& x : S) {
if (x < lo) lo = x;
if (x > hi) hi = x;
}
std::vector<compressed_value_type> buckets(hi - lo + 1, 0);
for (const auto& x : S) {
buckets[x - lo] = 1;
}
for (int v = 0; v < int(buckets.size()); v++) {
if (buckets[v]) buckets[v] = compressed_value_type(sigma++);
}
for (int i = 0; i < sz(S); i++) {
compressed_s[i] = buckets[S[i] - lo];
}
}
return construct_raw(compressed_s, sigma);
}
protected:
template <typename String> void build(const String& S, index_t sigma) {
N = sz(S);
build_sa(S, sigma);
build_rank();
}
private:
template <typename String> void build_sa(const String& S, index_t sigma) {
sa = std::vector<index_t>(N+1);
assert(sigma >= 0);
for (auto s : S) assert(0 <= index_t(s) && index_t(s) < sigma);
std::vector<index_t> tmp(sigma + std::max(N, sigma));
SuffixArrayBase::sais<String>(N, S, sa.data(), sigma, tmp.data());
}
template <typename String> static void sais(int N, const String& S, index_t* sa, int sigma, index_t* tmp) {
if (N == 0) {
sa[0] = 0;
return;
} else if (N == 1) {
sa[0] = 1;
sa[1] = 0;
return;
}
// Phase 1: Initialize the frequency array, which will let us lookup buckets.
index_t* freq = tmp; tmp += sigma;
memset(freq, 0, sizeof(*freq) * sigma);
for (int i = 0; i < N; i++) {
++freq[index_t(S[i])];
}
auto build_bucket_start = [&]() {
int cur = 1;
for (int v = 0; v < sigma; v++) {
tmp[v] = cur;
cur += freq[v];
}
};
auto build_bucket_end = [&]() {
int cur = 1;
for (int v = 0; v < sigma; v++) {
cur += freq[v];
tmp[v] = cur;
}
};
int num_pieces = 0;
int first_endpoint = 0;
// Phase 2: find the right-endpoints of the pieces
{
build_bucket_end();
// Initialize the final endpoint out-of-band this way so that we don't try to look up tmp[-1].
// This doesn't count towards num_pieces.
sa[0] = N;
index_t c0 = S[N-1], c1 = -1; bool isS = false;
for (int i = N-2; i >= 0; i--) {
c1 = c0;
c0 = S[i];
if (c0 < c1) {
isS = true;
} else if (c0 > c1 && isS) {
isS = false;
// insert i+1
sa[first_endpoint = --tmp[c1]] = i+1;
++num_pieces;
}
}
}
// If num_pieces <= 1, we don't need to actually run the recursion, it's just sorted automatically
// Otherwise, we're going to rebucket
if (num_pieces > 1) {
// Remove the first endpoint, we don't need to run the IS on this
sa[first_endpoint] = 0;
// Run IS for L-type
{
build_bucket_start();
for (int z = 0; z <= N; z++) {
int v = sa[z];
if (!v) continue;
// Leave for the S-round
if (v < 0) continue;
// clear out our garbage
sa[z] = 0;
--v;
index_t c0 = S[v-1], c1 = S[v];
sa[tmp[c1]++] = (c0 < c1) ? ~v : v;
}
}
index_t* const sa_end = sa + N + 1;
index_t* pieces = sa_end;
// Run IS for S-type and compactify
{
build_bucket_end();
for (int z = N; z >= 0; z--) {
int v = sa[z];
if (!v) continue;
// clear our garbage
sa[z] = 0;
if (v > 0) {
*--pieces = v;
continue;
}
v = ~v;
--v;
index_t c0 = S[v-1], c1 = S[v];
sa[--tmp[c1]] = (c0 > c1) ? v : ~v;
}
}
// Compute the lengths of the pieces in preparation for equality
// comparison, and store them in sa[v/2]. We set the length of the
// final piece to 0; it compares unequal to everything because of
// the sentinel.
{
int prv_start = N;
index_t c0 = S[N-1], c1 = -1; bool isS = false;
for (int i = N-2; i >= 0; i--) {
c1 = c0;
c0 = S[i];
if (c0 < c1) {
isS = true;
} else if (c0 > c1 && isS) {
isS = false;
// insert i+1
int v = i+1;
sa[v>>1] = prv_start == N ? 0 : prv_start - v;
prv_start = v;
}
}
}
// Compute the alphabet, storing the result into sa[v/2].
int next_sigma = 0;
{
int prv_len = -1, prv_v = 0;
for (int i = 0; i < num_pieces; i++) {
int v = pieces[i];
int len = sa[v>>1];
bool eq = prv_len == len;
for (int a = 0; eq && a < len; ++a) {
eq = S[v+a] == S[prv_v+a];
}
if (!eq) {
next_sigma++;
prv_len = len;
prv_v = v;
}
sa[v>>1] = next_sigma; // purposely leave this 1 large to check != 0
}
}
if (next_sigma == num_pieces) {
sa[0] = N;
memcpy(sa+1, pieces, sizeof(*sa) * num_pieces);
} else {
index_t* next_S = sa_end;
// Finally, pack the input to the SA
{
for (int i = (N-1)>>1; i >= 0; i--) {
int v = sa[i];
if (v) *--next_S = v-1;
sa[i] = 0;
}
}
memset(sa, 0, sizeof(*sa) * (num_pieces+1));
sais<const index_t*>(num_pieces, next_S, sa, next_sigma, tmp);
{ // Compute the piece start points again and use those to map up the suffix array
next_S = sa_end;
index_t c0 = S[N-1], c1 = -1; bool isS = false;
for (int i = N-2; i >= 0; i--) {
c1 = c0;
c0 = S[i];
if (c0 < c1) {
isS = true;
} else if (c0 > c1 && isS) {
isS = false;
int v = i+1;
*--next_S = v;
}
}
sa[0] = N;
for (int i = 1; i <= num_pieces; i++) {
sa[i] = next_S[sa[i]];
}
}
}
// zero everything else
memset(sa+num_pieces+1, 0, sizeof(*sa) * (N - num_pieces));
{
// Scatter the finished pieces
build_bucket_end();
for (int i = num_pieces; i > 0; i--) {
int v = sa[i];
sa[i] = 0;
index_t c1 = S[v];
sa[--tmp[c1]] = v;
}
}
}
// Home stretch! Just finish out with the L-type and then S-type
{
build_bucket_start();
for (int z = 0; z <= N; z++) {
int v = sa[z];
if (v <= 0) continue;
--v;
index_t c1 = S[v];
index_t c0 = v ? S[v-1] : c1; // if v = 0, we don't want to invert
sa[tmp[c1]++] = (c0 < c1) ? ~v : v;
}
}
// This just aggressively overwrites our original scattered pieces with the correct values
{
build_bucket_end();
for (int z = N; z >= 0; z--) {
int v = sa[z];
if (v >= 0) continue;
sa[z] = v = ~v;
--v;
index_t c1 = S[v];
index_t c0 = v ? S[v-1] : c1+1;
sa[--tmp[c1]] = (c0 > c1) ? v : ~v;
}
}
}
void build_rank() {
rank = std::vector<index_t>(N+1);
for (int i = 0; i <= N; i++) rank[sa[i]] = i;
}
};
class SuffixArray : public SuffixArrayBase<SuffixArray> {};
template <typename Self> class SuffixArrayLCPBase : public SuffixArrayBase<Self> {
public:
using index_t = typename SuffixArrayBase<Self>::index_t;
// lcp[i] = lcp(sa[i], sa[i+1])
std::vector<index_t> lcp;
protected:
friend SuffixArrayBase<Self>;
template <typename String> void build(const String& S, index_t sigma) {
SuffixArrayBase<Self>::build(S, sigma);
build_lcp(S);
}
private:
template <typename String> void build_lcp(const String& S) {
int N = this->N;
const auto& sa = this->sa;
const auto& rank = this->rank;
assert(sz(S) == N);
lcp = std::vector<index_t>(N);
for (int i = 0, k = 0; i < N - 1; i++) {
int j = sa[rank[i]-1];
while (k < N - std::max(i, j) && S[i+k] == S[j+k]) k++;
lcp[rank[i]-1] = k;
if (k) --k;
}
}
};
class SuffixArrayLCP : public SuffixArrayLCPBase<SuffixArrayLCP> {};
template <typename Self> class SuffixArrayRMQBase : public SuffixArrayLCPBase<Self> {
public:
using index_t = typename SuffixArrayLCPBase<Self>::index_t;
RangeMinQuery<std::pair<index_t, index_t>> rmq;
index_t get_lcp(index_t a, index_t b) const {
if (a == b) return this->N-a;
a = this->rank[a], b = this->rank[b];
if (a > b) std::swap(a, b);
return rmq.query(a, b-1).first;
}
// Get the split in the suffix tree, using half-open intervals
// Returns len, idx
std::pair<index_t, index_t> get_split(index_t l, index_t r) const {
assert(r - l > 1);
return rmq.query(l, r-2);
}
protected:
friend SuffixArrayBase<Self>;
template <typename String> void build(const String& S, index_t sigma) {
SuffixArrayLCPBase<Self>::build(S, sigma);
build_rmq();
}
private:
void build_rmq() {
int N = this->N;
const auto& lcp = this->lcp;
std::vector<std::pair<index_t, index_t>> lcp_idx(N);
for (int i = 0; i < N; i++) {
lcp_idx[i] = {lcp[i], i+1};
}
rmq = RangeMinQuery<std::pair<index_t, index_t>>(std::move(lcp_idx));
}
};
class SuffixArrayRMQ : public SuffixArrayRMQBase<SuffixArrayRMQ> {};
class PrefixArrayRMQ : private SuffixArrayRMQ {
PrefixArrayRMQ(const SuffixArrayRMQ& sa_) : SuffixArrayRMQ(sa_) {}
PrefixArrayRMQ(SuffixArrayRMQ&& sa_) : SuffixArrayRMQ(std::move(sa_)) {}
public:
PrefixArrayRMQ() {}
template <typename String> static PrefixArrayRMQ construct_raw(const String& S, int sigma) {
return PrefixArrayRMQ(SuffixArrayRMQ::construct_raw(String(S.rbegin(), S.rend()), sigma));
}
// TODO: Fill in other constructors
int get_lcs(int a, int b) const {
return SuffixArrayRMQ::get_lcp(N - a, N - b);
}
};
// clang-format off
// @formatter:off
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wpragmas"
#pragma GCC diagnostic ignored "-Wunknown-warning-option"
#pragma GCC diagnostic ignored "-Wmisleading-indentation"
#pragma GCC diagnostic ignored "-Wmultistatement-macros"
#include <bits/stdc++.h>
#include <cassert>
// src/rmq.hpp
template<typename T,class Compare=std::less<T>>class RangeMinQuery:private Compare{
static const int BUCKET_SIZE=32;
static const int BUCKET_SIZE_LOG=5;
static_assert(BUCKET_SIZE==(1<<BUCKET_SIZE_LOG),"BUCKET_SIZE should be a power of 2");
static const int CACHE_LINE_ALIGNMENT=64;
int n=0;
std::vector<T>data;
std::vector<T>pref_data;
std::vector<T>suff_data;
std::vector<T>sparse_table;
std::vector<uint32_t>range_mask;
private:
int num_buckets()const{
return n>>BUCKET_SIZE_LOG;
}
int num_levels()const{
return num_buckets()?32-__builtin_clz(num_buckets()):0;
}
int sparse_table_size()const{
return num_buckets()*num_levels();
}
private:
const T&min(const T&a,const T&b)const{
return Compare::operator()(a,b)?a:b;
}
void setmin(T&a,const T&b)const{
if(Compare::operator()(b,a))a=b;
}
template<typename Vec>static int get_size(const Vec&v){using std::size;return int(size(v));}
public:
RangeMinQuery(){}
template<typename Vec>explicit RangeMinQuery(const Vec&data_,const Compare&comp_=Compare())
:Compare(comp_)
,n(get_size(data_))
,data(n)
,pref_data(n)
,suff_data(n)
,sparse_table(sparse_table_size())
,range_mask(n)
{
for(int i=0;i<n;i++)data[i]=data_[i];
for(int i=0;i<n;i++){
if(i&(BUCKET_SIZE-1)){
uint32_t m=range_mask[i-1];
while(m&&!Compare::operator()(data[(i|(BUCKET_SIZE-1))-__builtin_clz(m)],data[i])){
m-=uint32_t(1)<<(BUCKET_SIZE-1-__builtin_clz(m));
}
m|=uint32_t(1)<<(i&(BUCKET_SIZE-1));
range_mask[i]=m;
}else{
range_mask[i]=1;
}
}
for(int i=0;i<n;i++){
pref_data[i]=data[i];
if(i&(BUCKET_SIZE-1)){
setmin(pref_data[i],pref_data[i-1]);
}
}
for(int i=n-1;i>=0;i--){
suff_data[i]=data[i];
if(i+1<n&&((i+1)&(BUCKET_SIZE-1))){
setmin(suff_data[i],suff_data[i+1]);
}
}
for(int i=0;i<num_buckets();i++){
sparse_table[i]=data[i*BUCKET_SIZE];
for(int v=1;v<BUCKET_SIZE;v++){
setmin(sparse_table[i],data[i*BUCKET_SIZE+v]);
}
}
for(int l=0;l+1<num_levels();l++){
for(int i=0;i+(1<<(l+1))<=num_buckets();i++){
sparse_table[(l+1)*num_buckets()+i]=min(sparse_table[l*num_buckets()+i],sparse_table[l*num_buckets()+i+(1<<l)]);
}
}
}
T query(int l,int r)const{
assert(l<=r);
int bucket_l=(l>>BUCKET_SIZE_LOG);
int bucket_r=(r>>BUCKET_SIZE_LOG);
if(bucket_l==bucket_r){
uint32_t msk=range_mask[r]&~((uint32_t(1)<<(l&(BUCKET_SIZE-1)))-1);
int ind=(l&~(BUCKET_SIZE-1))+__builtin_ctz(msk);
return data[ind];
}else{
T ans=min(suff_data[l],pref_data[r]);
bucket_l++;
if(bucket_l<bucket_r){
int level=(32-__builtin_clz(bucket_r-bucket_l))-1;
setmin(ans,sparse_table[level*num_buckets()+bucket_l]);
setmin(ans,sparse_table[level*num_buckets()+bucket_r-(1<<level)]);
}
return ans;
}
}
};
template<typename T>using RangeMaxQuery=RangeMinQuery<T,std::greater<T>>;
// src/suffix_array.hpp
template<class T>int sz(T&&arg){using std::size;return int(size(std::forward<T>(arg)));}
template<typename Self>class SuffixArrayBase{
public:
using index_t=int;
int N;
std::vector<index_t>sa;
std::vector<index_t>rank;
SuffixArrayBase():N(0){}
template<typename String>static Self construct_raw(const String&S,index_t sigma){
Self res;
res.build(S,sigma);
return res;
}
template<typename String,typename F>static Self map_and_construct(const String&S,const F&f,int sigma){
std::vector<decltype((f(S[0])))>mapped(sz(S));
for(int i=0;i<sz(S);i++){
mapped[i]=f(S[i]);
assert(0<=int(mapped[i])&&int(mapped[i])<sigma);
}
return construct_raw(mapped,sigma);
}
template<typename String>static Self sort_and_construct(const String&S){
using std::begin;
using std::end;
using value_type=typename std::iterator_traits<decltype(begin(S))>::value_type;
using compressed_value_type=typename std::conditional<
sizeof(value_type)<sizeof(index_t),
value_type,
index_t
>::type;
std::vector<compressed_value_type>compressed_s(sz(S));
int sigma=0;
{
std::vector<value_type>vals(begin(S),end(S));
std::sort(vals.begin(),vals.end());
vals.resize(unique(vals.begin(),vals.end())-vals.begin());
for(int i=0;i<sz(S);i++){
compressed_s[i]=compressed_value_type(index_t(std::lower_bound(vals.begin(),vals.end(),S[i])-vals.begin()));
}
sigma=int(vals.size());
}
return construct_raw(compressed_s,sigma);
}
template<typename String>static Self shift_and_construct(const String&S){
using std::begin;
using std::end;
using value_type=typename std::iterator_traits<decltype(begin(S))>::value_type;
std::vector<value_type>compressed_s(sz(S));
int sigma=0;
if(sz(S)>0){
value_type lo=*begin(S),hi=*begin(S);
for(const auto&x:S){
if(x<lo)lo=x;
if(x>hi)hi=x;
}
for(int i=0;i<sz(S);i++){
compressed_s[i]=value_type(S[i]-lo);
}
sigma=int(hi-lo+1);
}
return construct_raw(compressed_s,sigma);
}
template<typename String>static Self bucket_and_construct(const String&S){
using std::begin;
using std::end;
using value_type=typename std::iterator_traits<decltype(begin(S))>::value_type;
using compressed_value_type=typename std::conditional<
sizeof(value_type)<sizeof(index_t),
value_type,
index_t
>::type;
std::vector<compressed_value_type>compressed_s(sz(S));
int sigma=0;
if(sz(S)>0){
value_type lo=*begin(S),hi=*begin(S);
for(const auto&x:S){
if(x<lo)lo=x;
if(x>hi)hi=x;
}
std::vector<compressed_value_type>buckets(hi-lo+1,0);
for(const auto&x:S){
buckets[x-lo]=1;
}
for(int v=0;v<int(buckets.size());v++){
if(buckets[v])buckets[v]=compressed_value_type(sigma++);
}
for(int i=0;i<sz(S);i++){
compressed_s[i]=buckets[S[i]-lo];
}
}
return construct_raw(compressed_s,sigma);
}
protected:
template<typename String>void build(const String&S,index_t sigma){
N=sz(S);
build_sa(S,sigma);
build_rank();
}
private:
template<typename String>void build_sa(const String&S,index_t sigma){
sa=std::vector<index_t>(N+1);
assert(sigma>=0);
for(auto s:S)assert(0<=index_t(s)&&index_t(s)<sigma);
std::vector<index_t>tmp(sigma+std::max(N,sigma));
SuffixArrayBase::sais<String>(N,S,sa.data(),sigma,tmp.data());
}
template<typename String>static void sais(int N,const String&S,index_t*sa,int sigma,index_t*tmp){
if(N==0){
sa[0]=0;
return;
}else if(N==1){
sa[0]=1;
sa[1]=0;
return;
}
index_t*freq=tmp;tmp+=sigma;
memset(freq,0,sizeof(*freq)*sigma);
for(int i=0;i<N;i++){
++freq[index_t(S[i])];
}
auto build_bucket_start=[&](){
int cur=1;
for(int v=0;v<sigma;v++){
tmp[v]=cur;
cur+=freq[v];
}
};
auto build_bucket_end=[&](){
int cur=1;
for(int v=0;v<sigma;v++){
cur+=freq[v];
tmp[v]=cur;
}
};
int num_pieces=0;
int first_endpoint=0;
{
build_bucket_end();
sa[0]=N;
index_t c0=S[N-1],c1=-1;bool isS=false;
for(int i=N-2;i>=0;i--){
c1=c0;
c0=S[i];
if(c0<c1){
isS=true;
}else if(c0>c1&&isS){
isS=false;
sa[first_endpoint=--tmp[c1]]=i+1;
++num_pieces;
}
}
}
if(num_pieces>1){
sa[first_endpoint]=0;
{
build_bucket_start();
for(int z=0;z<=N;z++){
int v=sa[z];
if(!v)continue;
if(v<0)continue;
sa[z]=0;
--v;
index_t c0=S[v-1],c1=S[v];
sa[tmp[c1]++]=(c0<c1)?~v:v;
}
}
index_t*const sa_end=sa+N+1;
index_t*pieces=sa_end;
{
build_bucket_end();
for(int z=N;z>=0;z--){
int v=sa[z];
if(!v)continue;
sa[z]=0;
if(v>0){
*--pieces=v;
continue;
}
v=~v;
--v;
index_t c0=S[v-1],c1=S[v];
sa[--tmp[c1]]=(c0>c1)?v:~v;
}
}
{
int prv_start=N;
index_t c0=S[N-1],c1=-1;bool isS=false;
for(int i=N-2;i>=0;i--){
c1=c0;
c0=S[i];
if(c0<c1){
isS=true;
}else if(c0>c1&&isS){
isS=false;
int v=i+1;
sa[v>>1]=prv_start==N?0:prv_start-v;
prv_start=v;
}
}
}
int next_sigma=0;
{
int prv_len=-1,prv_v=0;
for(int i=0;i<num_pieces;i++){
int v=pieces[i];
int len=sa[v>>1];
bool eq=prv_len==len;
for(int a=0;eq&&a<len;++a){
eq=S[v+a]==S[prv_v+a];
}
if(!eq){
next_sigma++;
prv_len=len;
prv_v=v;
}
sa[v>>1]=next_sigma;
}
}
if(next_sigma==num_pieces){
sa[0]=N;
memcpy(sa+1,pieces,sizeof(*sa)*num_pieces);
}else{
index_t*next_S=sa_end;
{
for(int i=(N-1)>>1;i>=0;i--){
int v=sa[i];
if(v)*--next_S=v-1;
sa[i]=0;
}
}
memset(sa,0,sizeof(*sa)*(num_pieces+1));
sais<const index_t*>(num_pieces,next_S,sa,next_sigma,tmp);
{
next_S=sa_end;
index_t c0=S[N-1],c1=-1;bool isS=false;
for(int i=N-2;i>=0;i--){
c1=c0;
c0=S[i];
if(c0<c1){
isS=true;
}else if(c0>c1&&isS){
isS=false;
int v=i+1;
*--next_S=v;
}
}
sa[0]=N;
for(int i=1;i<=num_pieces;i++){
sa[i]=next_S[sa[i]];
}
}
}
memset(sa+num_pieces+1,0,sizeof(*sa)*(N-num_pieces));
{
build_bucket_end();
for(int i=num_pieces;i>0;i--){
int v=sa[i];
sa[i]=0;
index_t c1=S[v];
sa[--tmp[c1]]=v;
}
}
}
{
build_bucket_start();
for(int z=0;z<=N;z++){
int v=sa[z];
if(v<=0)continue;
--v;
index_t c1=S[v];
index_t c0=v?S[v-1]:c1;
sa[tmp[c1]++]=(c0<c1)?~v:v;
}
}
{
build_bucket_end();
for(int z=N;z>=0;z--){
int v=sa[z];
if(v>=0)continue;
sa[z]=v=~v;
--v;
index_t c1=S[v];
index_t c0=v?S[v-1]:c1+1;
sa[--tmp[c1]]=(c0>c1)?v:~v;
}
}
}
void build_rank(){
rank=std::vector<index_t>(N+1);
for(int i=0;i<=N;i++)rank[sa[i]]=i;
}
};
class SuffixArray:public SuffixArrayBase<SuffixArray>{};
template<typename Self>class SuffixArrayLCPBase:public SuffixArrayBase<Self>{
public:
using index_t=typename SuffixArrayBase<Self>::index_t;
std::vector<index_t>lcp;
protected:
friend SuffixArrayBase<Self>;
template<typename String>void build(const String&S,index_t sigma){
SuffixArrayBase<Self>::build(S,sigma);
build_lcp(S);
}
private:
template<typename String>void build_lcp(const String&S){
int N=this->N;
const auto&sa=this->sa;
const auto&rank=this->rank;
assert(sz(S)==N);
lcp=std::vector<index_t>(N);
for(int i=0,k=0;i<N-1;i++){
int j=sa[rank[i]-1];
while(k<N-std::max(i,j)&&S[i+k]==S[j+k])k++;
lcp[rank[i]-1]=k;
if(k)--k;
}
}
};
class SuffixArrayLCP:public SuffixArrayLCPBase<SuffixArrayLCP>{};
template<typename Self>class SuffixArrayRMQBase:public SuffixArrayLCPBase<Self>{
public:
using index_t=typename SuffixArrayLCPBase<Self>::index_t;
RangeMinQuery<std::pair<index_t,index_t>>rmq;
index_t get_lcp(index_t a,index_t b)const{
if(a==b)return this->N-a;
a=this->rank[a],b=this->rank[b];
if(a>b)std::swap(a,b);
return rmq.query(a,b-1).first;
}
std::pair<index_t,index_t>get_split(index_t l,index_t r)const{
assert(r-l>1);
return rmq.query(l,r-2);
}
protected:
friend SuffixArrayBase<Self>;
template<typename String>void build(const String&S,index_t sigma){
SuffixArrayLCPBase<Self>::build(S,sigma);
build_rmq();
}
private:
void build_rmq(){
int N=this->N;
const auto&lcp=this->lcp;
std::vector<std::pair<index_t,index_t>>lcp_idx(N);
for(int i=0;i<N;i++){
lcp_idx[i]={lcp[i],i+1};
}
rmq=RangeMinQuery<std::pair<index_t,index_t>>(std::move(lcp_idx));
}
};
class SuffixArrayRMQ:public SuffixArrayRMQBase<SuffixArrayRMQ>{};
class PrefixArrayRMQ:private SuffixArrayRMQ{
PrefixArrayRMQ(const SuffixArrayRMQ&sa_):SuffixArrayRMQ(sa_){}
PrefixArrayRMQ(SuffixArrayRMQ&&sa_):SuffixArrayRMQ(std::move(sa_)){}
public:
PrefixArrayRMQ(){}
template<typename String>static PrefixArrayRMQ construct_raw(const String&S,int sigma){
return PrefixArrayRMQ(SuffixArrayRMQ::construct_raw(String(S.rbegin(),S.rend()),sigma));
}
int get_lcs(int a,int b)const{
return SuffixArrayRMQ::get_lcp(N-a,N-b);
}
};
#pragma GCC diagnostic pop
// clang-format on
// @formatter:on
#pragma once
/*
* This is mostly inspired by https://golang.org/src/index/suffixarray/sais.go.
*/
#include <algorithm>
#include <vector>
#include <string>
#include <cassert>
#include <cstring>
#include <type_traits>
#include "rmq.hpp"
template<class T> int sz(T&& arg) { using std::size; return int(size(std::forward<T>(arg))); }
// Layered suffix array: SuffixArrayBase computes just sa/rank, each further
// layer statically opts into one more derived structure. Use the leaf classes
// SuffixArray, SuffixArrayLCP, or SuffixArrayRMQ; the named constructors on
// each return that type.
template <typename Self> class SuffixArrayBase {
public:
using index_t = int;
int N;
std::vector<index_t> sa;
std::vector<index_t> rank;
SuffixArrayBase() : N(0) {}
template <typename String> static Self construct_raw(const String& S, index_t sigma) {
Self res;
res.build(S, sigma);
return res;
}
// Pass a function which returns a value in [0, sigma)
template <typename String, typename F> static Self map_and_construct(const String& S, const F& f, int sigma) {
std::vector<decltype((f(S[0])))> mapped(sz(S));
for (int i = 0; i < sz(S); i++) {
mapped[i] = f(S[i]);
assert(0 <= int(mapped[i]) && int(mapped[i]) < sigma);
}
return construct_raw(mapped, sigma);
}
// Sorts the elements of S and then runs suffix array. This takes O(N log N) time with no dependence on sigma.
template <typename String> static Self sort_and_construct(const String& S) {
using std::begin;
using std::end;
using value_type = typename std::iterator_traits<decltype(begin(S))>::value_type;
using compressed_value_type = typename std::conditional<
sizeof(value_type) < sizeof(index_t),
value_type,
index_t
>::type;
std::vector<compressed_value_type> compressed_s(sz(S));
int sigma = 0;
{
std::vector<value_type> vals(begin(S), end(S));
std::sort(vals.begin(), vals.end());
vals.resize(unique(vals.begin(), vals.end()) - vals.begin());
for (int i = 0; i < sz(S); i++) {
compressed_s[i] = compressed_value_type(index_t(std::lower_bound(vals.begin(), vals.end(), S[i]) - vals.begin()));
}
sigma = int(vals.size());
}
return construct_raw(compressed_s, sigma);
}
// Shifts the elements so that sigma = max(S) - min(S) + 1
template <typename String> static Self shift_and_construct(const String& S) {
using std::begin;
using std::end;
using value_type = typename std::iterator_traits<decltype(begin(S))>::value_type;
std::vector<value_type> compressed_s(sz(S));
int sigma = 0;
if (sz(S) > 0) {
value_type lo = *begin(S), hi = *begin(S);
for (const auto& x : S) {
if (x < lo) lo = x;
if (x > hi) hi = x;
}
for (int i = 0; i < sz(S); i++) {
compressed_s[i] = value_type(S[i] - lo);
}
sigma = int(hi - lo + 1);
}
return construct_raw(compressed_s, sigma);
}
// Renumber/filter to only the used elements with bucket sorting. Still takes O(max(S) - min(S) + 1) memory/time,
// but should be less memory than `shift_and_construct` when sigma ~ N and max(S) - min(S) + 1 > N.
template <typename String> static Self bucket_and_construct(const String& S) {
using std::begin;
using std::end;
using value_type = typename std::iterator_traits<decltype(begin(S))>::value_type;
using compressed_value_type = typename std::conditional<
sizeof(value_type) < sizeof(index_t),
value_type,
index_t
>::type;
std::vector<compressed_value_type> compressed_s(sz(S));
int sigma = 0;
if (sz(S) > 0) {
value_type lo = *begin(S), hi = *begin(S);
for (const auto& x : S) {
if (x < lo) lo = x;
if (x > hi) hi = x;
}
std::vector<compressed_value_type> buckets(hi - lo + 1, 0);
for (const auto& x : S) {
buckets[x - lo] = 1;
}
for (int v = 0; v < int(buckets.size()); v++) {
if (buckets[v]) buckets[v] = compressed_value_type(sigma++);
}
for (int i = 0; i < sz(S); i++) {
compressed_s[i] = buckets[S[i] - lo];
}
}
return construct_raw(compressed_s, sigma);
}
protected:
template <typename String> void build(const String& S, index_t sigma) {
N = sz(S);
build_sa(S, sigma);
build_rank();
}
private:
template <typename String> void build_sa(const String& S, index_t sigma) {
sa = std::vector<index_t>(N+1);
assert(sigma >= 0);
for (auto s : S) assert(0 <= index_t(s) && index_t(s) < sigma);
std::vector<index_t> tmp(sigma + std::max(N, sigma));
SuffixArrayBase::sais<String>(N, S, sa.data(), sigma, tmp.data());
}
template <typename String> static void sais(int N, const String& S, index_t* sa, int sigma, index_t* tmp) {
if (N == 0) {
sa[0] = 0;
return;
} else if (N == 1) {
sa[0] = 1;
sa[1] = 0;
return;
}
// Phase 1: Initialize the frequency array, which will let us lookup buckets.
index_t* freq = tmp; tmp += sigma;
memset(freq, 0, sizeof(*freq) * sigma);
for (int i = 0; i < N; i++) {
++freq[index_t(S[i])];
}
auto build_bucket_start = [&]() {
int cur = 1;
for (int v = 0; v < sigma; v++) {
tmp[v] = cur;
cur += freq[v];
}
};
auto build_bucket_end = [&]() {
int cur = 1;
for (int v = 0; v < sigma; v++) {
cur += freq[v];
tmp[v] = cur;
}
};
int num_pieces = 0;
int first_endpoint = 0;
// Phase 2: find the right-endpoints of the pieces
{
build_bucket_end();
// Initialize the final endpoint out-of-band this way so that we don't try to look up tmp[-1].
// This doesn't count towards num_pieces.
sa[0] = N;
index_t c0 = S[N-1], c1 = -1; bool isS = false;
for (int i = N-2; i >= 0; i--) {
c1 = c0;
c0 = S[i];
if (c0 < c1) {
isS = true;
} else if (c0 > c1 && isS) {
isS = false;
// insert i+1
sa[first_endpoint = --tmp[c1]] = i+1;
++num_pieces;
}
}
}
// If num_pieces <= 1, we don't need to actually run the recursion, it's just sorted automatically
// Otherwise, we're going to rebucket
if (num_pieces > 1) {
// Remove the first endpoint, we don't need to run the IS on this
sa[first_endpoint] = 0;
// Run IS for L-type
{
build_bucket_start();
for (int z = 0; z <= N; z++) {
int v = sa[z];
if (!v) continue;
// Leave for the S-round
if (v < 0) continue;
// clear out our garbage
sa[z] = 0;
--v;
index_t c0 = S[v-1], c1 = S[v];
sa[tmp[c1]++] = (c0 < c1) ? ~v : v;
}
}
index_t* const sa_end = sa + N + 1;
index_t* pieces = sa_end;
// Run IS for S-type and compactify
{
build_bucket_end();
for (int z = N; z >= 0; z--) {
int v = sa[z];
if (!v) continue;
// clear our garbage
sa[z] = 0;
if (v > 0) {
*--pieces = v;
continue;
}
v = ~v;
--v;
index_t c0 = S[v-1], c1 = S[v];
sa[--tmp[c1]] = (c0 > c1) ? v : ~v;
}
}
// Compute the lengths of the pieces in preparation for equality
// comparison, and store them in sa[v/2]. We set the length of the
// final piece to 0; it compares unequal to everything because of
// the sentinel.
{
int prv_start = N;
index_t c0 = S[N-1], c1 = -1; bool isS = false;
for (int i = N-2; i >= 0; i--) {
c1 = c0;
c0 = S[i];
if (c0 < c1) {
isS = true;
} else if (c0 > c1 && isS) {
isS = false;
// insert i+1
int v = i+1;
sa[v>>1] = prv_start == N ? 0 : prv_start - v;
prv_start = v;
}
}
}
// Compute the alphabet, storing the result into sa[v/2].
int next_sigma = 0;
{
int prv_len = -1, prv_v = 0;
for (int i = 0; i < num_pieces; i++) {
int v = pieces[i];
int len = sa[v>>1];
bool eq = prv_len == len;
for (int a = 0; eq && a < len; ++a) {
eq = S[v+a] == S[prv_v+a];
}
if (!eq) {
next_sigma++;
prv_len = len;
prv_v = v;
}
sa[v>>1] = next_sigma; // purposely leave this 1 large to check != 0
}
}
if (next_sigma == num_pieces) {
sa[0] = N;
memcpy(sa+1, pieces, sizeof(*sa) * num_pieces);
} else {
index_t* next_S = sa_end;
// Finally, pack the input to the SA
{
for (int i = (N-1)>>1; i >= 0; i--) {
int v = sa[i];
if (v) *--next_S = v-1;
sa[i] = 0;
}
}
memset(sa, 0, sizeof(*sa) * (num_pieces+1));
sais<const index_t*>(num_pieces, next_S, sa, next_sigma, tmp);
{ // Compute the piece start points again and use those to map up the suffix array
next_S = sa_end;
index_t c0 = S[N-1], c1 = -1; bool isS = false;
for (int i = N-2; i >= 0; i--) {
c1 = c0;
c0 = S[i];
if (c0 < c1) {
isS = true;
} else if (c0 > c1 && isS) {
isS = false;
int v = i+1;
*--next_S = v;
}
}
sa[0] = N;
for (int i = 1; i <= num_pieces; i++) {
sa[i] = next_S[sa[i]];
}
}
}
// zero everything else
memset(sa+num_pieces+1, 0, sizeof(*sa) * (N - num_pieces));
{
// Scatter the finished pieces
build_bucket_end();
for (int i = num_pieces; i > 0; i--) {
int v = sa[i];
sa[i] = 0;
index_t c1 = S[v];
sa[--tmp[c1]] = v;
}
}
}
// Home stretch! Just finish out with the L-type and then S-type
{
build_bucket_start();
for (int z = 0; z <= N; z++) {
int v = sa[z];
if (v <= 0) continue;
--v;
index_t c1 = S[v];
index_t c0 = v ? S[v-1] : c1; // if v = 0, we don't want to invert
sa[tmp[c1]++] = (c0 < c1) ? ~v : v;
}
}
// This just aggressively overwrites our original scattered pieces with the correct values
{
build_bucket_end();
for (int z = N; z >= 0; z--) {
int v = sa[z];
if (v >= 0) continue;
sa[z] = v = ~v;
--v;
index_t c1 = S[v];
index_t c0 = v ? S[v-1] : c1+1;
sa[--tmp[c1]] = (c0 > c1) ? v : ~v;
}
}
}
void build_rank() {
rank = std::vector<index_t>(N+1);
for (int i = 0; i <= N; i++) rank[sa[i]] = i;
}
};
class SuffixArray : public SuffixArrayBase<SuffixArray> {};
template <typename Self> class SuffixArrayLCPBase : public SuffixArrayBase<Self> {
public:
using index_t = typename SuffixArrayBase<Self>::index_t;
// lcp[i] = lcp(sa[i], sa[i+1])
std::vector<index_t> lcp;
protected:
friend SuffixArrayBase<Self>;
template <typename String> void build(const String& S, index_t sigma) {
SuffixArrayBase<Self>::build(S, sigma);
build_lcp(S);
}
private:
template <typename String> void build_lcp(const String& S) {
int N = this->N;
const auto& sa = this->sa;
const auto& rank = this->rank;
assert(sz(S) == N);
lcp = std::vector<index_t>(N);
for (int i = 0, k = 0; i < N - 1; i++) {
int j = sa[rank[i]-1];
while (k < N - std::max(i, j) && S[i+k] == S[j+k]) k++;
lcp[rank[i]-1] = k;
if (k) --k;
}
}
};
class SuffixArrayLCP : public SuffixArrayLCPBase<SuffixArrayLCP> {};
template <typename Self> class SuffixArrayRMQBase : public SuffixArrayLCPBase<Self> {
public:
using index_t = typename SuffixArrayLCPBase<Self>::index_t;
RangeMinQuery<std::pair<index_t, index_t>> rmq;
index_t get_lcp(index_t a, index_t b) const {
if (a == b) return this->N-a;
a = this->rank[a], b = this->rank[b];
if (a > b) std::swap(a, b);
return rmq.query(a, b-1).first;
}
// Get the split in the suffix tree, using half-open intervals
// Returns len, idx
std::pair<index_t, index_t> get_split(index_t l, index_t r) const {
assert(r - l > 1);
return rmq.query(l, r-2);
}
protected:
friend SuffixArrayBase<Self>;
template <typename String> void build(const String& S, index_t sigma) {
SuffixArrayLCPBase<Self>::build(S, sigma);
build_rmq();
}
private:
void build_rmq() {
int N = this->N;
const auto& lcp = this->lcp;
std::vector<std::pair<index_t, index_t>> lcp_idx(N);
for (int i = 0; i < N; i++) {
lcp_idx[i] = {lcp[i], i+1};
}
rmq = RangeMinQuery<std::pair<index_t, index_t>>(std::move(lcp_idx));
}
};
class SuffixArrayRMQ : public SuffixArrayRMQBase<SuffixArrayRMQ> {};
class PrefixArrayRMQ : private SuffixArrayRMQ {
PrefixArrayRMQ(const SuffixArrayRMQ& sa_) : SuffixArrayRMQ(sa_) {}
PrefixArrayRMQ(SuffixArrayRMQ&& sa_) : SuffixArrayRMQ(std::move(sa_)) {}
public:
PrefixArrayRMQ() {}
template <typename String> static PrefixArrayRMQ construct_raw(const String& S, int sigma) {
return PrefixArrayRMQ(SuffixArrayRMQ::construct_raw(String(S.rbegin(), S.rend()), sigma));
}
// TODO: Fill in other constructors
int get_lcs(int a, int b) const {
return SuffixArrayRMQ::get_lcp(N - a, N - b);
}
};