New Distinct Substrings 【 后缀数组 模板题 】

Input:
2
CCCCC
ABABA

Output:
5
9

#include<bits/stdc++.h>

using namespace std;
const int maxn = 200010;
int cntA[maxn],cntB[maxn],tsa[maxn];// cntA 、cntB 、 tsa 都是中间的变量。
int A[maxn],B[maxn];
int sa[maxn]; // sa[i] 代表第i小的后缀位置。
int Rank[maxn]; // Rank[i] 代表位置 i 开始的后缀是第几小，也就是排名。
int height[maxn]; // 把所有的后缀串按字典序从小到大排序， height[i]表示第i小的后缀串与第i-1小的后缀串的最长公共前缀是多少
long long n;//串的长度的平方会超过int
char ch[maxn];
void solve()//求sa[],Rank[],height[]
{
    for(int i = 0; i < 256; i++) cntA[i] = 0; //初始化
    for(int i = 1; i <= n; i++) cntA[ch[i-1]]++; // 统计字符出现的次数
    for(int i = 1; i < 256; i++) cntA[i] += cntA[i-1];//求前缀和
    for(int i = n; i; i--) sa[cntA[ch[i-1]]--] = i;
    Rank[sa[1]] = 1; // 第1小的位置也就是“这个位置是第1小”
    for(int i = 2; i <= n; i++)  // 这个for就是把Rank[]求出来
    {
        Rank[sa[i]] = Rank[sa[i-1]]; // 相连接的Rank相同，因为比较的时候满足长度相同，下面也一样
        if(ch[sa[i]-1] != ch[sa[i-1]-1]) Rank[sa[i]]++; // 如果不相同，就需要增加了
    }
    // 以上求长度为1的后缀的Rank[]，就是每个字符对应的Rank[]

    for(int l = 1; Rank[sa[n]] < n; l <<= 1) // 倍增来求Rank[]
    {
        memset(cntA, 0, sizeof(cntA));
        memset(cntB, 0, sizeof(cntB));
        for(int i = 1; i <= n; i++)
        {
            cntA[A[i] = Rank[i]]++;
            cntB[B[i] = (i+l <= n)?Rank[i+l]:0]++; // 位数不够需要补零
        }
        for(int i = 1; i <= n; i++) cntB[i] += cntB[i-1];
        for(int i = n; i; i--) tsa[cntB[B[i]]--] = i;
        for(int i = 1; i <= n; i++) cntA[i] += cntA[i-1];
        for(int i = n; i; i--) sa[cntA[A[tsa[i]]]--] = tsa[i];
        //上面为了求sa[]数组。i是第一类，i+l是第二类，tsa[i]代表第二类第i小的后缀位置

        Rank[sa[1]]=1; // 和上面差不多
        for(int i = 2; i <= n; i++)
        {
            Rank[sa[i]] = Rank[sa[i-1]];
            if(A[sa[i]] != A[sa[i-1]] || B[sa[i]] != B[sa[i-1]]) Rank[sa[i]]++; // 这里就需要看论文了
        }
    }
    //求height[]，任意两个后缀串的最长公共前缀 = 它们之间 height[] 数组的最小值。
    for(int i = 1, j = 0; i <= n; i++)
    {
        if(j) j--;
        while(ch[i+j-1] == ch[sa[Rank[i]-1] + j - 1]) j++;
        height[Rank[i]] = j;
    }
}
int main()
{
    int T;
    scanf("%d", &T);
    while(T--)
    {
        scanf("%s", ch);
        n = strlen(ch);
        if(n == 1)
        {
            printf("1\n");
            continue;
        }
        solve();
        long long ans = (n+1)*n/2;
        for(int i = 2; i <= n; i++) ans -= height[i];//减去所有重复的字串
        printf("%lld\n", ans);
    }
}