test/yosupo/number_of_substrings.test.cpp

View this file on GitHub
Last update: 2024-01-07 22:37:45+09:00
Problem: https://judge.yosupo.jp/problem/number_of_substrings

Depends on

Code

#define PROBLEM "https://judge.yosupo.jp/problem/number_of_substrings"

#include <bits/stdc++.h>

#include "../../string/lcp_array.hpp"
#include "../../string/suffix_array.hpp"
using namespace std;

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(0);

    string S;
    cin >> S;
    int N = S.size();
    auto sa = suffix_array(S);
    auto lcp = lcp_array(S, sa);
    long long ans = 1LL * N * (N + 1) / 2;
    for (int i = 0; i < lcp.size(); ++i) ans -= lcp[i];
    cout << ans << endl;
}

#line 1 "test/yosupo/number_of_substrings.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/number_of_substrings"

#include <bits/stdc++.h>

#line 4 "string/lcp_array.hpp"

template <typename T>
std::vector<int> lcp_array(const std::vector<T>& s,
                           const std::vector<int>& sa) {
    const int n = s.size();
    std::vector<int> rank(n);
    for (int i = 0; i < n; ++i) rank[sa[i]] = i;
    int h = 0;
    std::vector<int> lcp(n - 1);
    for (int i = 0; i < n; ++i) {
        if (h > 0) --h;
        if (rank[i] == 0) continue;
        int j = sa[rank[i] - 1];
        while (j + h < n && i + h < n && s[j + h] == s[i + h]) ++h;
        lcp[rank[i] - 1] = h;
    }
    return lcp;
}

std::vector<int> lcp_array(const std::string& s, const std::vector<int>& sa) {
    return lcp_array(std::vector<char>(s.begin(), s.end()), sa);
}
#line 6 "string/suffix_array.hpp"

template <typename T>
std::vector<int> suffix_array(const std::vector<T>& s) {
    const int n = s.size();
    std::vector<int> sa(n);
    std::iota(sa.begin(), sa.end(), 0);
    std::ranges::sort(sa, {}, [&](int i) { return s[i]; });
    int cl = 0;
    std::vector<int> rank(n);
    for (int i = 1; i < n; ++i) {
        if (s[sa[i - 1]] != s[sa[i]]) ++cl;
        rank[sa[i]] = cl;
    }
    std::vector<int> tmp(n), nrank(n), cnt(n);
    for (int k = 1; k < n; k <<= 1) {
        // sort by second half
        int cnt1 = 0, cnt2 = k;
        for (int i = 0; i < n; ++i) {
            int j = sa[i] - k;
            if (j >= 0)
                tmp[cnt2++] = j;
            else
                tmp[cnt1++] = j + n;
        }

        // sort by first half
        std::ranges::fill(cnt, 0);
        for (int i = 0; i < n; ++i) ++cnt[rank[tmp[i]]];
        for (int i = 1; i < n; ++i) cnt[i] += cnt[i - 1];
        for (int i = n - 1; i >= 0; --i) sa[--cnt[rank[tmp[i]]]] = tmp[i];

        // assign new rank
        nrank[sa[0]] = 0;
        cl = 0;
        for (int i = 1; i < n; ++i) {
            if (rank[sa[i - 1]] != rank[sa[i]] ||
                (sa[i - 1] + k < n ? rank[sa[i - 1] + k] : -1) !=
                    (sa[i] + k < n ? rank[sa[i] + k] : -1)) {
                ++cl;
            }
            nrank[sa[i]] = cl;
        }
        rank.swap(nrank);
    }
    return sa;
}

std::vector<int> suffix_array(const std::string& s) {
    return suffix_array(std::vector<char>(s.begin(), s.end()));
}

template <typename T>
std::vector<int> cyclic_suffix_array(const std::vector<T>& s) {
    const int n = s.size();
    std::vector<int> sa(n);
    std::iota(sa.begin(), sa.end(), 0);
    std::ranges::sort(sa, {}, [&](int i) { return s[i]; });
    int cl = 0;
    std::vector<int> rank(n);
    for (int i = 1; i < n; ++i) {
        if (s[sa[i - 1]] != s[sa[i]]) ++cl;
        rank[sa[i]] = cl;
    }
    std::vector<int> tmp(n), nrank(n), cnt(n);
    for (int k = 1; k < n; k <<= 1) {
        // sort by second half
        int cnt1 = 0, cnt2 = k;
        for (int i = 0; i < n; ++i) {
            int j = sa[i] - k;
            if (j >= 0)
                tmp[cnt2++] = j;
            else
                tmp[cnt1++] = j + n;
        }

        // sort by first half
        std::ranges::fill(cnt, 0);
        for (int i = 0; i < n; ++i) ++cnt[rank[tmp[i]]];
        for (int i = 1; i < n; ++i) cnt[i] += cnt[i - 1];
        for (int i = n - 1; i >= 0; --i) sa[--cnt[rank[tmp[i]]]] = tmp[i];

        // assign new rank
        nrank[sa[0]] = 0;
        cl = 0;
        for (int i = 1; i < n; ++i) {
            if (rank[sa[i - 1]] != rank[sa[i]] ||
                rank[(sa[i - 1] + k) % n] != rank[(sa[i] + k) % n]) {
                ++cl;
            }
            nrank[sa[i]] = cl;
        }
        rank.swap(nrank);
    }
    return sa;
}

std::vector<int> cyclic_suffix_array(const std::string& s) {
    return cyclic_suffix_array(std::vector<char>(s.begin(), s.end()));
}

/*
// comparator for substrings
// used for string matching with the suffix array

auto cmp = [&](int si, const string& t) {
    int sn = S.size(), tn = t.size();
    int ti = 0;
    for (; si < sn && ti < tn; ++si, ++ti) {
        if (T[si] < t[ti]) return true;
        if (T[si] > t[ti]) return false;
    }
    return si == sn && ti < tn;
};
*/
#line 7 "test/yosupo/number_of_substrings.test.cpp"
using namespace std;

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(0);

    string S;
    cin >> S;
    int N = S.size();
    auto sa = suffix_array(S);
    auto lcp = lcp_array(S, sa);
    long long ans = 1LL * N * (N + 1) / 2;
    for (int i = 0; i < lcp.size(); ++i) ans -= lcp[i];
    cout << ans << endl;
}