test/yosupo/number_of_substrings.SA.test.cpp

• View this file on GitHub
• Last update: 2024-09-01 12:11:12+09:00
• Problem: https://judge.yosupo.jp/problem/number_of_substrings

Depends on

• Suffix Array (接尾辞配列) (src/String/SuffixArray.hpp)

Code

// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/number_of_substrings
// competitive-verifier: TLE 0.5
// competitive-verifier: MLE 64
#include <iostream>
#include <string>
#include "src/String/SuffixArray.hpp"
using namespace std;
signed main() {
 cin.tie(0);
 ios::sync_with_stdio(0);
 string S;
 cin >> S;
 SuffixArray sa(S);
 LCPArray lcp(sa);
 int N= S.length();
 long long ans= (long long)N * (N + 1) / 2;
 for (int i= N; --i;) ans-= lcp[i - 1];
 cout << ans << '\n';
 return 0;
}

#line 1 "test/yosupo/number_of_substrings.SA.test.cpp"
// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/number_of_substrings
// competitive-verifier: TLE 0.5
// competitive-verifier: MLE 64
#include <iostream>
#include <string>
#line 3 "src/String/SuffixArray.hpp"
#include <vector>
#include <algorithm>
template <class String> struct SuffixArray {
 String s;
 std::vector<int> sa;
 static inline std::vector<int> sa_is(const std::vector<int> &s, int K) {
  const int n= s.size();
  std::vector<char> t(n);
  std::vector<int> bkt(K, 0), bkt_l(K), bkt_r(K), sa(n), p1;
  t.back()= true;
  for (int i= n; --i;)
   if (t[i - 1]= (s[i - 1] < s[i] || (t[i] && s[i - 1] == s[i])); t[i] && !t[i - 1]) p1.push_back(i);
  std::reverse(p1.begin(), p1.end());
  const int n1= p1.size();
  for (int i= n; i--;) ++bkt[s[i]];
  for (int i= 0, sum= 0; i < K; ++i) sum+= bkt[i], bkt_r[i]= sum, bkt_l[i]= sum - bkt[i];
  std::vector<int> s1(n1), sa1(n1);
  std::fill_n(sa.begin(), n, -1), std::copy_n(bkt_r.begin(), K, bkt.begin());
  for (int i= n1; i--;) sa[--bkt[s[p1[i]]]]= p1[i];
  std::copy_n(bkt_l.begin(), K, bkt.begin());
  for (int i= 0, j; i < n; ++i)
   if ((j= sa[i] - 1) >= 0 && !t[j]) sa[bkt[s[j]]++]= j;
  std::copy_n(bkt_r.begin(), K, bkt.begin());
  for (int i= n, j; i--;)
   if ((j= sa[i] - 1) >= 0 && t[j]) sa[--bkt[s[j]]]= j;
  for (int i= 0, j= 0; i < n; ++i)
   if (t[sa[i]] && sa[i] > 0 && !t[sa[i] - 1]) sa1[j++]= sa[i];
  int name= 0;
  for (int i= 0, prev= -1, j, pos; i < n1; ++i, sa[pos]= name - 1)
   for (j= 0, pos= sa1[i];; ++j)
    if (prev == -1 || s[pos + j] != s[prev + j] || t[pos + j] != t[prev + j]) {
     prev= pos, ++name;
     break;
    } else if (j && ((t[pos + j] && !t[pos + j - 1]) || (t[prev + j] && !t[prev + j - 1]))) break;
  for (int i= n1; i--;) s1[i]= sa[p1[i]];
  if (name != n1) sa1= sa_is(s1, name);
  else
   for (int i= n1; i--;) sa1[s1[i]]= i;
  std::copy_n(bkt_r.begin(), K, bkt.begin()), std::fill_n(sa.begin(), n, -1);
  for (int i= n1; i--;) sa[--bkt[s[p1[sa1[i]]]]]= p1[sa1[i]];
  for (int i= 0, j; i < n; ++i)
   if ((j= sa[i] - 1) >= 0 && !t[j]) sa[bkt_l[s[j]]++]= j;
  for (int i= n, j; i--;)
   if ((j= sa[i] - 1) >= 0 && t[j]) sa[--bkt_r[s[j]]]= j;
  return sa;
 }
public:
 SuffixArray(const String &S): s(S) {
  std::vector<int> s_cpy(s.size() + 1);
  if constexpr (std::is_convertible_v<String, std::string>) std::copy(s.begin(), s.end(), s_cpy.begin()), sa= sa_is(s_cpy, 128), sa.erase(sa.begin());
  else {
   auto v= s;
   sort(v.begin(), v.end()), v.erase(unique(v.begin(), v.end()), v.end());
   for (int i= s.size(); i--;) s_cpy[i]= std::lower_bound(v.begin(), v.end(), s[i]) - v.begin() + 1;
   sa= sa_is(s_cpy, v.size() + 1), sa.erase(sa.begin());
  }
 }
 int operator[](int i) const { return sa[i]; }
 size_t size() const { return sa.size(); }
 auto begin() const { return sa.begin(); }
 auto end() const { return sa.end(); }
 // return {l,r} s.t. P is a prefix of S[sa[i]:] ( i in [l,r) )
 // l == r if P is not a substring of S
 // O(|P|log|S|)
 std::pair<int, int> pattern_matching(const String &P) const {
  const int n= s.size(), m= P.size();
  if (n < m) return {0, 0};
  auto f1= [&](int h) {
   auto t= s.begin() + h;
   for (int j= 0, e= std::min(n - h, m); j < e; ++j) {
    if (t[j] < P[j]) return true;
    if (t[j] > P[j]) return false;
   }
   return n - h < m;
  };
  auto f2= [&](int h) {
   auto t= s.begin() + h;
   for (int j= 0, e= std::min(n - h, m); j < e; ++j)
    if (t[j] > P[j]) return false;
   return true;
  };
  auto L= std::partition_point(sa.begin(), sa.end(), f1), R= std::partition_point(L, sa.end(), f2);
  return {L - sa.begin(), R - sa.begin()};
 }
};
struct LCPArray {
 std::vector<int> rnk;
 template <class String> LCPArray(const SuffixArray<String> &sa): rnk(sa.size()) {
  const int n= sa.size(), log= n > 2 ? 31 - __builtin_clz(n - 2) : 0;
  dat.resize(log + 1), dat[0].resize(n - 1);
  auto &lcp= dat[0];
  for (int i= n; i--;) rnk[sa[i]]= i;
  for (int i= 0, h= 0; i < n; ++i) {
   if (rnk[i] == n - 1) {
    h= 0;
    continue;
   }
   for (int j= sa[rnk[i] + 1]; i + h < n && j + h < n && sa.s[i + h] == sa.s[j + h];) ++h;
   if ((lcp[rnk[i]]= h)) --h;
  }
  for (int i= 0, I= 1, j; i < log; ++i, I<<= 1)
   for (dat[i + 1].resize(j= dat[i].size() - I); j--;) dat[i + 1][j]= std::min(dat[i][j], dat[i][j + I]);
 }
 int operator[](int i) const { return dat[0][i]; }
 size_t size() const { return dat[0].size(); }
 auto begin() const { return dat[0].begin(); }
 auto end() const { return dat[0].end(); }
 int operator()(int i, int j) const {
  if (i == j) return rnk.size() - i;
  auto [l, r]= std::minmax(rnk[i], rnk[j]);
  if (r == l + 1) return dat[0][l];
  int k= 31 - __builtin_clz(r - l - 1);
  return std::min(dat[k][l], dat[k][r - (1 << k)]);
 }
private:
 std::vector<std::vector<int>> dat;
};
#line 7 "test/yosupo/number_of_substrings.SA.test.cpp"
using namespace std;
signed main() {
 cin.tie(0);
 ios::sync_with_stdio(0);
 string S;
 cin >> S;
 SuffixArray sa(S);
 LCPArray lcp(sa);
 int N= S.length();
 long long ans= (long long)N * (N + 1) / 2;
 for (int i= N; --i;) ans-= lcp[i - 1];
 cout << ans << '\n';
 return 0;
}

Test cases

Env	Name	Status	Elapsed	Memory
g++-13	all_same_00	AC	37 ms	43 MB
g++-13	all_same_01	AC	36 ms	43 MB
g++-13	all_same_02	AC	36 ms	43 MB
g++-13	all_same_03	AC	35 ms	43 MB
g++-13	all_same_04	AC	35 ms	43 MB
g++-13	example_00	AC	5 ms	4 MB
g++-13	example_01	AC	4 ms	3 MB
g++-13	example_02	AC	4 ms	3 MB
g++-13	example_03	AC	4 ms	4 MB
g++-13	fib_str_00	AC	45 ms	44 MB
g++-13	fib_str_01	AC	34 ms	33 MB
g++-13	fib_str_02	AC	36 ms	31 MB
g++-13	fib_str_03	AC	29 ms	29 MB
g++-13	fib_str_04	AC	46 ms	43 MB
g++-13	max_random_00	AC	59 ms	43 MB
g++-13	max_random_01	AC	60 ms	43 MB
g++-13	max_random_02	AC	59 ms	43 MB
g++-13	max_random_03	AC	60 ms	43 MB
g++-13	max_random_04	AC	59 ms	42 MB
g++-13	random_00	AC	47 ms	34 MB
g++-13	random_01	AC	55 ms	40 MB
g++-13	random_02	AC	10 ms	7 MB
g++-13	random_03	AC	51 ms	37 MB
g++-13	random_04	AC	34 ms	25 MB
clang++-18	all_same_00	AC	32 ms	43 MB
clang++-18	all_same_01	AC	33 ms	43 MB
clang++-18	all_same_02	AC	33 ms	43 MB
clang++-18	all_same_03	AC	32 ms	43 MB
clang++-18	all_same_04	AC	32 ms	43 MB
clang++-18	example_00	AC	5 ms	4 MB
clang++-18	example_01	AC	5 ms	4 MB
clang++-18	example_02	AC	4 ms	4 MB
clang++-18	example_03	AC	4 ms	4 MB
clang++-18	fib_str_00	AC	41 ms	44 MB
clang++-18	fib_str_01	AC	32 ms	33 MB
clang++-18	fib_str_02	AC	32 ms	31 MB
clang++-18	fib_str_03	AC	28 ms	29 MB
clang++-18	fib_str_04	AC	43 ms	43 MB
clang++-18	max_random_00	AC	56 ms	43 MB
clang++-18	max_random_01	AC	57 ms	43 MB
clang++-18	max_random_02	AC	60 ms	43 MB
clang++-18	max_random_03	AC	56 ms	43 MB
clang++-18	max_random_04	AC	56 ms	42 MB
clang++-18	random_00	AC	45 ms	34 MB
clang++-18	random_01	AC	53 ms	40 MB
clang++-18	random_02	AC	10 ms	7 MB
clang++-18	random_03	AC	49 ms	37 MB
clang++-18	random_04	AC	32 ms	25 MB

Back to top page