test/string/Number_of_Substrings.test.cpp

View this file on GitHub
Last update: 2023-11-17 01:51:05+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 "../../template/template.hpp"
#include "../../string/suffix_array.hpp"

namespace lib {

void main_() {
    string s;
    cin >> s;
    auto sa = suffix_array(s);
    auto lcp = lcp_array(s, sa);
    ll n = s.size();
    ll ans = n * (n + 1) / 2;
    rep(i,0,n-1) {
        ans -= lcp[i];
    }
    cout << ans << '\n';
}

}  // namespace ebi

int main() {
    int t = 1;
    // cin >> t;
    while (t--) {
        lib::main_();
    }
    return 0;
}

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

#line 2 "template/template.hpp"

#include <bits/stdc++.h>

#define rep(i, s, n) for (int i = (int)(s); i < (int)(n); i++)
#define rrep(i, s, n) for (int i = (int)(n)-1; i >= (int)(s); i--)
#define all(v) v.begin(), v.end()

using ll = long long;
using ld = long double;
using ull = unsigned long long;

template <typename T> bool chmin(T &a, const T &b) {
    if (a <= b) return false;
    a = b;
    return true;
}
template <typename T> bool chmax(T &a, const T &b) {
    if (a >= b) return false;
    a = b;
    return true;
}

namespace lib {

using namespace std;

}  // namespace lib

// using namespace lib;
#line 2 "string/suffix_array.hpp"

#line 4 "string/suffix_array.hpp"

namespace lib {

vector<int> suffix_array(const string &s) {
    int n = s.size();
    vector<int> rnk(n, -1), sa(n);
    iota(all(sa), 0);
    rep(i, 0, n) {
        rnk[i] = s[i];
    }
    int k = 1;
    auto compare = [&](int i, int j) -> bool {
        if (rnk[i] != rnk[j]) return rnk[i] < rnk[j];
        return (i + k < n ? rnk[i + k] : -1) < (j + k < n ? rnk[j + k] : -1);
    };
    while (k < n) {
        sort(all(sa), compare);
        vector<int> tmp(n, 0);
        rep(i, 1, n) {
            tmp[sa[i]] = tmp[sa[i - 1]] + int(compare(sa[i - 1], sa[i]));
        }
        swap(rnk, tmp);
        k <<= 1;
    }
    return sa;
}

vector<int> lcp_array(const string &s, const vector<int> &sa) {
    int n = s.size();
    vector<int> rnk(n), lcp(n - 1);
    rep(i, 0, n) rnk[sa[i]] = i;
    int h = 0;
    rep(i, 0, n) {
        if (h > 0) h--;
        if (rnk[i] == 0) continue;
        int j = sa[rnk[i] - 1];
        while (i + h < n && j + h < n && s[i + h] == s[j + h]) {
            h++;
        }
        lcp[rnk[i] - 1] = h;
    }
    return lcp;
}

}  // namespace lib
#line 5 "test/string/Number_of_Substrings.test.cpp"

namespace lib {

void main_() {
    string s;
    cin >> s;
    auto sa = suffix_array(s);
    auto lcp = lcp_array(s, sa);
    ll n = s.size();
    ll ans = n * (n + 1) / 2;
    rep(i,0,n-1) {
        ans -= lcp[i];
    }
    cout << ans << '\n';
}

}  // namespace ebi

int main() {
    int t = 1;
    // cin >> t;
    while (t--) {
        lib::main_();
    }
    return 0;
}