Codeforces 802I Fake News (hard)

Codeforces 802I

题意：统计所有不同子串出现次数的平方的和。

想法：建一个SAM，$Ans=\sum (step[i]-step[fa[i]])*right[i]^2$

#include<cstdio>
#include<cstring>
#include<algorithm>

typedef long long ll;
template
inline void read(T&x)
{
	x=0;bool f=0;char c=getchar();
	while((c<'0'||c>'9')&&c!='-')c=getchar(); if(c=='-')f=1,c=getchar();
	while(c>='0'&&c<='9'){x=x*10+c-'0';c=getchar();}
	x=f?-x:x;
}
const int MAXN(100010);
struct SAM{int nx[26],step,pre,right;}sam[MAXN<<1];int lastson,last,now=1,top=1,root=1;
int T,len;ll Ans;
char str[MAXN];
void extend(int x)
{
	last=now; now=++top; sam[now].step=sam[last].step+1; sam[now].right=1;
	for(;!sam[last].nx[x]&&last;last=sam[last].pre)
		sam[last].nx[x]=now;
	if(!last)sam[now].pre=root;
	else
	{
		lastson=sam[last].nx[x];
		if(sam[lastson].step==sam[last].step+1)sam[now].pre=lastson;
		else 
		{
			sam[++top]=sam[lastson];sam[top].right=0;
			sam[top].step=sam[last].step+1;
			sam[lastson].pre=sam[now].pre=top;
			for(;sam[last].nx[x]==lastson&&last;last=sam[last].pre)
				sam[last].nx[x]=top;
		}
	}
}
int cnt[MAXN],p[MAXN<<1];
void Total()
{
	for(int i=1;i<=len;i++)cnt[i]=0;
	for(int i=1;i<=top;i++)cnt[sam[i].step]++;
	for(int i=1;i<=len;i++)cnt[i]+=cnt[i-1];
	for(int i=top;i>=1;i--)p[cnt[sam[i].step]--]=i;
	for(int i=top;i>=1;i--)sam[sam[p[i]].pre].right+=sam[p[i]].right;
	for(int i=1;i<=top;i++)
	{
		Ans+=(ll)sam[i].right*sam[i].right*(sam[i].step-sam[sam[i].pre].step);
//		fprintf(stderr,"%lld %d\n",(ll)sam[i].right*sam[i].right,sam[i].step);
	}
}
int main()
{
//	freopen("C.in","r",stdin);
	read(T);
	while(T--)
	{
		memset(sam,0,sizeof(sam));top=root=now=1;Ans=0;
		scanf("%s",str+1);len=strlen(str+1);
		for(int i=1;i<=len;i++)extend(str[i]-'a');
		Total();
		printf("%lld\n",Ans);
	}
	return 0;
}