cp-book

ecnerwala's competitive programming library

View the Project on GitHub ecnerwala/cp-book

:heavy_check_mark: verify/number_of_substrings.test.cpp

View this file on GitHub · Last update: 2026-07-25 03:32:41-07:00

Problem: https://judge.yosupo.jp/problem/number_of_substrings

Depends on

Code

// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/number_of_substrings

#include <bits/stdc++.h>
#include <cassert>

#include "suffix_array.hpp"

int main() {
	std::ios_base::sync_with_stdio(false), std::cin.tie(nullptr);

	std::string S; std::cin >> S;
	int N = int(S.size());
	auto sa = SuffixArrayLCP::shift_and_construct(S);
	int64_t ans = int64_t(N) * int64_t(N+1) / 2;
	for (int i = 0; i < N; i++) ans -= sa.lcp[i];
	std::cout << ans << '\n';

	return 0;
}
#include <bits/stdc++.h>
#line 1 "verify/number_of_substrings.test.cpp"
// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/number_of_substrings

#line 5 "verify/number_of_substrings.test.cpp"

#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);
	}
};
#line 7 "verify/number_of_substrings.test.cpp"

int main() {
	std::ios_base::sync_with_stdio(false), std::cin.tie(nullptr);

	std::string S; std::cin >> S;
	int N = int(S.size());
	auto sa = SuffixArrayLCP::shift_and_construct(S);
	int64_t ans = int64_t(N) * int64_t(N+1) / 2;
	for (int i = 0; i < N; i++) ans -= sa.lcp[i];
	std::cout << ans << '\n';

	return 0;
}
// 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>
// 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);
}
};
// verify/number_of_substrings.test.cpp
int main(){
std::ios_base::sync_with_stdio(false),std::cin.tie(nullptr);
std::string S;std::cin>>S;
int N=int(S.size());
auto sa=SuffixArrayLCP::shift_and_construct(S);
int64_t ans=int64_t(N)*int64_t(N+1)/2;
for(int i=0;i<N;i++)ans-=sa.lcp[i];
std::cout<<ans<<'\n';
return 0;
}
#pragma GCC diagnostic pop
// clang-format on
// @formatter:on

Test cases

Env Name Status Elapsed Memory
g++-sanitizer all_same_00 :heavy_check_mark: AC 50 ms 19 MB
g++-sanitizer all_same_01 :heavy_check_mark: AC 38 ms 20 MB
g++-sanitizer all_same_02 :heavy_check_mark: AC 41 ms 19 MB
g++-sanitizer all_same_03 :heavy_check_mark: AC 47 ms 19 MB
g++-sanitizer all_same_04 :heavy_check_mark: AC 40 ms 19 MB
g++-sanitizer example_00 :heavy_check_mark: AC 16 ms 8 MB
g++-sanitizer example_01 :heavy_check_mark: AC 14 ms 9 MB
g++-sanitizer example_02 :heavy_check_mark: AC 17 ms 8 MB
g++-sanitizer example_03 :heavy_check_mark: AC 16 ms 9 MB
g++-sanitizer fib_str_00 :heavy_check_mark: AC 70 ms 20 MB
g++-sanitizer fib_str_01 :heavy_check_mark: AC 56 ms 18 MB
g++-sanitizer fib_str_02 :heavy_check_mark: AC 56 ms 18 MB
g++-sanitizer fib_str_03 :heavy_check_mark: AC 58 ms 16 MB
g++-sanitizer fib_str_04 :heavy_check_mark: AC 76 ms 21 MB
g++-sanitizer max_random_00 :heavy_check_mark: AC 89 ms 21 MB
g++-sanitizer max_random_01 :heavy_check_mark: AC 88 ms 21 MB
g++-sanitizer max_random_02 :heavy_check_mark: AC 96 ms 21 MB
g++-sanitizer max_random_03 :heavy_check_mark: AC 91 ms 21 MB
g++-sanitizer max_random_04 :heavy_check_mark: AC 89 ms 21 MB
g++-sanitizer random_00 :heavy_check_mark: AC 80 ms 19 MB
g++-sanitizer random_01 :heavy_check_mark: AC 85 ms 20 MB
g++-sanitizer random_02 :heavy_check_mark: AC 21 ms 10 MB
g++-sanitizer random_03 :heavy_check_mark: AC 93 ms 20 MB
g++-sanitizer random_04 :heavy_check_mark: AC 58 ms 16 MB
g++ all_same_00 :heavy_check_mark: AC 14 ms 10 MB
g++ all_same_01 :heavy_check_mark: AC 10 ms 10 MB
g++ all_same_02 :heavy_check_mark: AC 9 ms 10 MB
g++ all_same_03 :heavy_check_mark: AC 9 ms 10 MB
g++ all_same_04 :heavy_check_mark: AC 10 ms 10 MB
g++ example_00 :heavy_check_mark: AC 2 ms 4 MB
g++ example_01 :heavy_check_mark: AC 2 ms 4 MB
g++ example_02 :heavy_check_mark: AC 2 ms 4 MB
g++ example_03 :heavy_check_mark: AC 2 ms 4 MB
g++ fib_str_00 :heavy_check_mark: AC 19 ms 10 MB
g++ fib_str_01 :heavy_check_mark: AC 14 ms 8 MB
g++ fib_str_02 :heavy_check_mark: AC 15 ms 8 MB
g++ fib_str_03 :heavy_check_mark: AC 12 ms 8 MB
g++ fib_str_04 :heavy_check_mark: AC 19 ms 10 MB
g++ max_random_00 :heavy_check_mark: AC 39 ms 10 MB
g++ max_random_01 :heavy_check_mark: AC 39 ms 10 MB
g++ max_random_02 :heavy_check_mark: AC 39 ms 10 MB
g++ max_random_03 :heavy_check_mark: AC 39 ms 10 MB
g++ max_random_04 :heavy_check_mark: AC 38 ms 10 MB
g++ random_00 :heavy_check_mark: AC 30 ms 9 MB
g++ random_01 :heavy_check_mark: AC 36 ms 10 MB
g++ random_02 :heavy_check_mark: AC 5 ms 4 MB
g++ random_03 :heavy_check_mark: AC 33 ms 9 MB
g++ random_04 :heavy_check_mark: AC 22 ms 7 MB
Back to top page