Test/yosupo-Number-of-Substrings.test.cpp

View this file on GitHub
Last update: 2023-07-19 16:58:05+00:00
Problem: https://judge.yosupo.jp/problem/number_of_substrings

Depends on

Code

#include<bits/stdc++.h>
using namespace std;

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

#include "../String/SuffixArray.hpp"

#include "../String/LCPArray.hpp"

int main() {

    string S;
    cin >> S;

    vector<int> lcp = LCPArray(S, SuffixArray(S));

    long long ans = (long long)S.size() * (long long)(S.size() + 1) / 2;
    for (int v : lcp) {
        ans -= v;
    }

    cout << ans << endl;

}

#line 1 "Test/yosupo-Number-of-Substrings.test.cpp"
#include<bits/stdc++.h>
using namespace std;

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

#line 1 "String/SuffixArray.hpp"
vector<int> SuffixArray(const string& s) {
    int n = s.size(), k;
    vector<int> sa(n + 1);
    vector<int> rank(n + 1);
    for (int i = 0; i <= n; i++) {
        sa[i] = i;
        rank[i] = i < n ? (int)s[i] : -1;
    }

    auto compare = [&](int a, int b) {
        if (rank[a] != rank[b]) return rank[a] < rank[b];
        auto rank_ak = a + k <= n ? rank[a + k] : -1;
        auto rank_bk = b + k <= n ? rank[b + k] : -1;
        return rank_ak < rank_bk;
    };

    for (k = 1; k <= n; k *= 2) {
        sort(sa.begin(), sa.end(), compare);
        vector<int> tmp(n + 1);
        tmp[sa[0]] = 0;
        for (int i = 1; i <= n; i++) {
            tmp[sa[i]] = tmp[sa[i - 1]] + (compare(sa[i - 1], sa[i]) ? 1 : 0);
        }
        for (int i = 0; i <= n; i++) {
            rank[i] = tmp[i];
        }
    }
    return sa;
}
#line 7 "Test/yosupo-Number-of-Substrings.test.cpp"

#line 1 "String/LCPArray.hpp"
vector<int> LCPArray(const string& s, const vector<int>& sa) {
    int n = s.size();
    vector<int> rank(n + 1);
    for (int i = 0; i <= n; i++) {
        rank[sa[i]] = i;
    }
    vector<int> lcp(n, 0);
    int h = 0;
    for (int i = 0; i < n; i++) {
        int j = sa[rank[i] - 1];
        if (h > 0) h--;
        for (; j + h < n && i + h < n; h++) {
            if (s[j + h] != s[i + h]) break;
        }
        lcp[rank[i] - 1] = h;
    }
    return lcp;
}
#line 9 "Test/yosupo-Number-of-Substrings.test.cpp"

int main() {

    string S;
    cin >> S;

    vector<int> lcp = LCPArray(S, SuffixArray(S));

    long long ans = (long long)S.size() * (long long)(S.size() + 1) / 2;
    for (int v : lcp) {
        ans -= v;
    }

    cout << ans << endl;

}