BZOJ3238: [Ahoi2013]差异 （后缀自动机）

Description

#include<set>
#include<map>
#include<ctime>
#include<queue>
#include<cmath>
#include<cstdio>
#include<vector>
#include<cstring>
#include<cstdlib>
#include<iostream>
#include<algorithm>
#define N 500005
#define inf 1000000000
#define pa pair<int,int>
#define ll long long
using namespace std;
ll ans;
int n,k,p,q=1,top;
int v
,a
,h
,sa[2]
,rk[2]
;
int st
,l
,r
;
char ch
;
void mul(int *sa,int *rk,int *SA,int *RK)
{
for(int i=1;i<=n;i++)v[rk[sa[i]]]=i;
for(int i=n;i;i--)
if(sa[i]>k)
SA[v[rk[sa[i]-k]]--]=sa[i]-k;
for(int i=n-k+1;i<=n;i++)SA[v[rk[i]]--]=i;
for(int i=1;i<=n;i++)
RK[SA[i]]=RK[SA[i-1]]+(rk[SA[i-1]]!=rk[SA[i]]||rk[SA[i-1]+k]!=rk[SA[i]+k]);
}
void presa()
{
for(int i=1;i<=n;i++)v[a[i]]++;
for(int i=1;i<=30;i++)v[i]+=v[i-1];
for(int i=1;i<=n;i++)sa[p][v[a[i]]--]=i;
for(int i=1;i<=n;i++)
rk[p][sa[p][i]]=rk[p][sa[p][i-1]]+(a[sa[p][i-1]]!=a[sa[p][i]]);
for(k=1;k<n;k<<=1,swap(p,q))
mul(sa[p],rk[p],sa[q],rk[q]);
for(int k=0,i=1;i<=n;i++)
{
int j=sa[p][rk[p][i]-1];
while(ch[j+k]==ch[i+k])k++;
h[rk[p][i]]=k;if(k>0)k--;
}
}
void solve()
{
for(int i=1;i<=n;i++)ans+=(ll)i*(n-1);
h[0]=-inf;
for(int i=1;i<=n;i++)
{
while(h[i]<=h[st[top]])top--;
if(st[top]==0)l[i]=1;
else l[i]=st[top]+1;
st[++top]=i;
}
h[n+1]=-inf;top=0;st[0]=n+1;
for(int i=n;i;i--)
{
while(h[i]<h[st[top]])top--;
if(st[top]==n+1)r[i]=n;
else r[i]=st[top]-1;
st[++top]=i;
}
for(int i=1;i<=n;i++)
ans-=2LL*(i-l[i]+1)*(r[i]-i+1)*h[i];
}
int main()
{
scanf("%s",ch+1);
n=strlen(ch+1);
for(int i=1;i<=n;i++)a[i]=ch[i]-'a'+1;
presa();
solve();
printf("%lld",ans);
return 0;
}