test/string/suffix_array.test.cpp

• View this file on GitHub
• Last update: 2024-07-11 22:21:57+09:00
• Problem: https://judge.yosupo.jp/problem/suffixarray

Code

#define PROBLEM "https://judge.yosupo.jp/problem/suffixarray"
#include "include/mtl/string/suffix_array.hpp"
#include <iostream>
using namespace std;

int main() {
  string s; cin>>s;
  auto sa = SuffixArray(s.begin(), s.end());
  for (auto v : sa) cout << v << ' ';
  cout << endl;
}

#line 1 "test/string/suffix_array.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/suffixarray"
#line 2 "include/mtl/string/suffix_array.hpp"
#include <vector>
#include <iterator>
#include <limits>
#include <cassert>
#include <algorithm>
#include <numeric>
#include <cstring>
#include <iostream>

void sa_is(const std::vector<int>& S, std::vector<std::pair<int, int>>& bucket, std::vector<int>::iterator dst, bool definitely_distinct_lmss) {
  const int n = (int) S.size();
  assert(S.back() == 0);
  assert(*max_element(S.begin(), S.end()) < n);
  // Induced Suffix Sorting
  if (n <= 3) {
    if (n == 1)
      dst[0] = 0;
    else if (n == 2) {
      dst[0] = 1;
      dst[1] = 0;
    } else if (n == 3) {
      dst[0] = 2;
      if (S[0] < S[1]) {
        dst[1] = 0;
        dst[2] = 1;
      } else {
        dst[1] = 1;
        dst[2] = 0;
      }
    }
    return;
  }
  std::vector<bool> is_s(n);
  is_s[n-1] = true;
  for (int i = n-3; i >= 0; i--) {
    int l = S[i], r = S[i+1];
    is_s[i] = l < r or (l == r and is_s[i+1]);
  }
  auto is_lms = [&](int i) {
    return i and is_s[i] and !is_s[i-1];
  };

  std::vector<int> cnt(n);
  for (int c : S)
    cnt[c]++;
  assert(bucket.size() >= (size_t)n);
  auto init_bucket = [&bucket, n](const std::vector<int>& cnt) {
    int sum = 0;
    for (int i = 0; i < n; i++) {
      bucket[i] = {sum, sum+cnt[i]-1};
      sum += cnt[i];
    }
  };
  init_bucket(cnt);

  auto A = &*dst;
  memset(A, -1, sizeof(int)*n);
  int cnt_lms = 0;
  for (int i = n-1; i >= 1; i--) if (is_lms(i)) {
      A[bucket[S[i]].second--] = i;
      cnt_lms++;
  }
  auto induced_sort = [&]() {
    for (int i = 0; i < n; i++) {
      if (A[i] == -1) continue;
      int p = A[i] ? A[i]-1 : n-1;
      if (!is_s[p])
        A[bucket[S[p]].first++] = p;
      if (is_lms(A[i])) {
        assert(bucket[S[A[i]]].second+1 == i);
        A[++bucket[S[A[i]]].second] = -1;
      }
    }
    for (int i = n-1; i >= 0; i--) {
      assert(A[i] != -1);
      int p = A[i] ? A[i]-1 : n-1;
      if (is_s[p])
        A[bucket[S[p]].second--] = p;
    }
  };
  induced_sort();

  if (definitely_distinct_lmss)
    return;

  // Build SA of LMS-substrings
  std::vector<int> sa_lmss(cnt_lms), lmss(n);
  {
    auto are_same_lms_substr = [&](int i, int j) {
      assert(is_lms(i) and is_lms(j));
      if (S[i] != S[j])
        return false;
      i++; j++;
      while(true) {
        if (S[i] != S[j] or is_lms(i) != is_lms(j))
          return false;
        if (is_lms(i))
          break;
        i++; j++;
      }
      return true;
    };
    int lms_ord = 0;
    int plms = n-1;
    auto cmp = [&](int i) -> int& {
      return lmss[cnt_lms + i/2];
    };
    assert(A[0] == n-1);
    cmp(A[0]) = 0;
    for (int i = 1; i < n; i++) if (is_lms(A[i])) {
      if (!are_same_lms_substr(plms, A[i]))
        ++lms_ord;
      cmp(A[i]) = lms_ord;
      plms = A[i];
    }
    int k = 0;
    for (int i = 1; i < n; i++) if (is_lms(i))
      lmss[k++] = cmp(i);
    lmss.resize(cnt_lms);
    sa_is(lmss, bucket, sa_lmss.begin(), cnt_lms - 1 == lms_ord);
  }

  // Induced-sort by ordered LMSs
  auto& citi = lmss; // Reuse memory
  citi.clear();
  for (int i = 1; i < n; i++) if (is_lms(i))
    citi.push_back(i);
  init_bucket(cnt);
  memset(A, -1, sizeof(int)*n);
  for (auto it = sa_lmss.rbegin(); it != sa_lmss.rend(); ++it)
    A[bucket[S[citi[*it]]].second--] = citi[*it];
  induced_sort();
}

template<typename It>
[[nodiscard]]
std::vector<int> SuffixArray(It begin, It end) {
  using trait = std::iterator_traits<It>;
  static_assert(std::is_base_of<std::random_access_iterator_tag, typename trait::iterator_category>::value, "");

  int n = end - begin;
  std::vector<int> S(n+1);
  S[n] = 0;
  auto mp = std::minmax_element(begin, end);
  if (0 < (int)*mp.first and (int)*mp.second <= n) {
    std::copy(begin, end, S.begin());
  } else {
    std::vector<int> ord(n);
    std::iota(ord.begin(), ord.end(), 0);
    std::sort(ord.begin(), ord.end(), [&](int l, int r) {
      return *(begin+l) < *(begin+r);
    });
    int k = 1;
    for (int i = 0; i < n; i++)
      S[ord[i]] = i and *(begin+ord[i]) != *(begin+ord[i-1]) ? ++k : k;
  }
  std::vector<int> A(n+1);
  std::vector<std::pair<int,int>> bucket(n+1);
  sa_is(S, bucket, A.begin(), false);
  return std::vector<int>(A.begin()+1, A.end());
}
#line 4 "test/string/suffix_array.test.cpp"
using namespace std;

int main() {
  string s; cin>>s;
  auto sa = SuffixArray(s.begin(), s.end());
  for (auto v : sa) cout << v << ' ';
  cout << endl;
}

Test cases

Env	Name	Status	Elapsed	Memory
g++	all_same_00	AC	61 ms	15 MB
g++	all_same_01	AC	60 ms	16 MB
g++	all_same_02	AC	64 ms	15 MB
g++	all_same_03	AC	59 ms	16 MB
g++	all_same_04	AC	59 ms	15 MB
g++	binary_carry_00	AC	97 ms	21 MB
g++	binary_carry_01	AC	93 ms	21 MB
g++	example_00	AC	6 ms	3 MB
g++	example_01	AC	5 ms	3 MB
g++	example_02	AC	5 ms	3 MB
g++	example_03	AC	5 ms	3 MB
g++	fib_str_00	AC	80 ms	19 MB
g++	fib_str_01	AC	60 ms	15 MB
g++	fib_str_02	AC	63 ms	15 MB
g++	fib_str_03	AC	50 ms	12 MB
g++	fib_str_04	AC	81 ms	19 MB
g++	hack_00	AC	5 ms	3 MB
g++	hack_01	AC	5 ms	3 MB
g++	hack_02	AC	5 ms	3 MB
g++	max_random_00	AC	92 ms	18 MB
g++	max_random_01	AC	100 ms	18 MB
g++	max_random_02	AC	101 ms	18 MB
g++	max_random_03	AC	100 ms	18 MB
g++	max_random_04	AC	91 ms	18 MB
g++	near_power_of_2_max_random_00	AC	51 ms	11 MB
g++	near_power_of_2_max_random_01	AC	51 ms	11 MB
g++	near_power_of_2_max_same_00	AC	34 ms	10 MB
g++	near_power_of_2_max_same_01	AC	37 ms	10 MB
g++	one_00	AC	5 ms	3 MB
g++	random_00	AC	74 ms	15 MB
g++	random_01	AC	91 ms	17 MB
g++	random_02	AC	14 ms	5 MB
g++	random_03	AC	82 ms	16 MB
g++	random_04	AC	54 ms	11 MB
g++	small_random_00	AC	5 ms	3 MB
g++	small_random_01	AC	5 ms	4 MB
g++	small_random_02	AC	5 ms	3 MB
g++	small_random_03	AC	5 ms	3 MB
g++	small_random_04	AC	5 ms	3 MB
g++	small_random_05	AC	5 ms	3 MB
g++	small_random_06	AC	5 ms	3 MB
g++	small_random_07	AC	5 ms	3 MB
g++	small_random_08	AC	5 ms	3 MB
g++	small_random_09	AC	5 ms	4 MB
clang++	all_same_00	AC	56 ms	15 MB
clang++	all_same_01	AC	55 ms	15 MB
clang++	all_same_02	AC	55 ms	15 MB
clang++	all_same_03	AC	64 ms	16 MB
clang++	all_same_04	AC	64 ms	15 MB
clang++	binary_carry_00	AC	93 ms	21 MB
clang++	binary_carry_01	AC	92 ms	21 MB
clang++	example_00	AC	5 ms	3 MB
clang++	example_01	AC	5 ms	3 MB
clang++	example_02	AC	5 ms	3 MB
clang++	example_03	AC	5 ms	3 MB
clang++	fib_str_00	AC	81 ms	19 MB
clang++	fib_str_01	AC	58 ms	15 MB
clang++	fib_str_02	AC	64 ms	15 MB
clang++	fib_str_03	AC	51 ms	13 MB
clang++	fib_str_04	AC	82 ms	19 MB
clang++	hack_00	AC	5 ms	3 MB
clang++	hack_01	AC	5 ms	3 MB
clang++	hack_02	AC	5 ms	3 MB
clang++	max_random_00	AC	92 ms	18 MB
clang++	max_random_01	AC	90 ms	18 MB
clang++	max_random_02	AC	91 ms	18 MB
clang++	max_random_03	AC	85 ms	18 MB
clang++	max_random_04	AC	85 ms	18 MB
clang++	near_power_of_2_max_random_00	AC	52 ms	11 MB
clang++	near_power_of_2_max_random_01	AC	49 ms	11 MB
clang++	near_power_of_2_max_same_00	AC	32 ms	10 MB
clang++	near_power_of_2_max_same_01	AC	31 ms	10 MB
clang++	one_00	AC	5 ms	3 MB
clang++	random_00	AC	69 ms	15 MB
clang++	random_01	AC	80 ms	17 MB
clang++	random_02	AC	14 ms	5 MB
clang++	random_03	AC	74 ms	16 MB
clang++	random_04	AC	55 ms	11 MB
clang++	small_random_00	AC	5 ms	3 MB
clang++	small_random_01	AC	5 ms	4 MB
clang++	small_random_02	AC	5 ms	3 MB
clang++	small_random_03	AC	5 ms	4 MB
clang++	small_random_04	AC	5 ms	3 MB
clang++	small_random_05	AC	5 ms	4 MB
clang++	small_random_06	AC	5 ms	3 MB
clang++	small_random_07	AC	5 ms	3 MB
clang++	small_random_08	AC	5 ms	3 MB
clang++	small_random_09	AC	5 ms	3 MB

Back to top page