bzoj 3238 ahoi2013差異 後綴數組

#include<set>
#include<cmath>
#include<cstdio>
#include<cstring>
#include<iostream>
#include<algorithm>
#define maxn 500005
#define LL long long
using namespace std;
char s[maxn];
set<int> S;
set<int> :: iterator it;
struct Suffix{
    int n;
    int sa[maxn],rank[maxn],height[maxn];
    int cnt[maxn],sec[maxn],tmp[maxn];
    void get_sa()
    {
        n=strlen(s+1);
        for(int i=1;i<=n;i++) rank[i]=s[i]-'a'+1;
        for(int i=1;i<=n;i++) cnt[rank[i]]++;
        for(int i=1;i<=26;i++) cnt[i]+=cnt[i-1];
        for(int i=1;i<=n;i++) sa[cnt[rank[i]]--]=i;

        for(int k=1;k<=n;k<<=1)
        {
            int top=0;
            for(int i=n-k+1;i<=n;i++) sec[++top]=i;
            for(int i=1;i<=n;i++) if(sa[i]>k) sec[++top]=sa[i]-k;

            for(int i=0;i<=n;i++) cnt[i]=0;
            for(int i=1;i<=n;i++) cnt[rank[i]]++;
            for(int i=1;i<=n;i++) cnt[i]+=cnt[i-1];
            for(int i=n;i>=1;i--) sa[cnt[rank[sec[i]]]--]=sec[i];

            top=tmp[sa[1]]=1;
            for(int i=2;i<=n;i++)
            {
                if(rank[sa[i]]!=rank[sa[i-1]]||rank[sa[i]+k]!=rank[sa[i-1]+k]) top++;
                tmp[sa[i]]=top;
            }
            for(int i=1;i<=n;i++) rank[i]=tmp[i];
        }
        int k=0;
        for(int i=1;i<=n;i++)
        {
            if(k) k--;
            while(i+k<=n&&sa[rank[i]-1]+k<=n&&s[i+k]==s[sa[rank[i]-1]+k]) k++;
            height[rank[i]]=k;
        }
    }
}SA;
bool cmp(int A,int B)
{return SA.height[A]<SA.height[B];}
int pos[maxn];
int main()
{
    scanf("%s",s+1);
    int n=strlen(s+1);
    LL ans=(LL)(1+n)*(LL)n/2ll;ans*=(LL)(n-1);
    SA.get_sa();
    for(int i=1;i<n;i++) pos[i]=i+1;
    sort(pos+1,pos+n,cmp);
    S.insert(1);S.insert(n+1);
    int l,r;
    for(int i=1;i<n;i++)
    {
        it=S.lower_bound(pos[i]);it--;l=*it;
        it=S.lower_bound(pos[i]);r=*it;
        ans-=(LL)2*(LL)SA.height[pos[i]]*(LL)(r-pos[i])*(LL)(pos[i]-l);
        S.insert(pos[i]);
    }
    cout<<ans;
    return 0;
}