广义后缀自动机例题

P3346 [ZJOI2015]诸神眷顾的幻想乡

以每个叶子节点为开头进行dfs遍历，将遍历到的串全部加入建立广义SAM，结果即为本质不同的字符串个数

#include<cstdio>
#include<cstring>
#include<algorithm>
#include<iostream>
#include<queue>
#define int long long
using namespace std;
const int maxn=2e6+10;
int n,m,tot,len[maxn],fa[maxn],ch[maxn][26],c;
int head[maxn<<1],nex[maxn<<1],to[maxn<<1];
int cnt[maxn];
char s[maxn];
int vis[maxn];
int res;
int col[maxn];
int k=1;
int ecnt;
void add(int x,int y) {
	to[++ecnt]=y;
	nex[ecnt]=head[x];
	head[x]=ecnt;
}
int Ins(int c,int last) {
	int p=last;
	if(ch[p][c]) {
		int q=ch[p][c];
		if(len[p]+1==len[q]) {
			return q;
		} else {
			int nq=++tot;
			len[nq]=len[p]+1;
			memcpy(ch[nq],ch[q],sizeof(ch[nq]));
			fa[nq]=fa[q];
			fa[q]=nq;
			for(; p&&ch[p][c]==q; p=fa[p])ch[p][c]=nq;
			return nq;
		}
	}
	int np=++tot;
	len[np]=len[p]+1;
	for(; p&&!ch[p][c]; p=fa[p])ch[p][c]=np;
	if(!p)fa[np]=1;
	else {
		int q=ch[p][c];
		if(len[p]+1==len[q])fa[np]=q;
		else {
			int nq=++tot;
			vis[nq]=k;
			len[nq]=len[p]+1;
			memcpy(ch[nq],ch[q],sizeof(ch[nq]));
			fa[nq]=fa[q];
			fa[q]=fa[np]=nq;
			for(; p&&ch[p][c]==q; p=fa[p])ch[p][c]=nq;
		}
	}
	return np;
}
int id[maxn];
inline void prepare() {
	for (int i=1; i<=tot; i++) cnt[len[i]]++;
	for (int i=1; i<=tot; i++) cnt[i]+=cnt[i-1];
	for (int i=1; i<=tot; i++) id[cnt[len[i]]--]=i;
	for (int i=tot; i>=1; i--) {
		int X=id[i];
		if (vis[fa[X]]==0) vis[fa[X]]=vis[X];
		else if (vis[fa[X]]!=vis[X]) vis[fa[X]]=-1;
	}
}
void dfs(int x,int fa,int last) {
	int tmp=Ins(col[x],last);
	for(int i=head[x]; i; i=nex[i]) {
		int y=to[i];
		if(y==fa)continue;
		dfs(y,x,tmp);
	}
}
int in[maxn];
signed main() {
	std::ios::sync_with_stdio(false);
	int n,c;
	cin>>n>>c;
	tot=1;
	for(int i=1; i<=n; i++) {
		cin>>col[i];
	}
	for(int i=1; i<n; i++) {
		int x,y;
		cin>>x>>y;
		in[x]++;
		in[y]++;
		add(x,y);
		add(y,x);
	}
	for(int i=1; i<=n; i++) {
		if(in[i]==1)
		dfs(i,i,1);
	}
	for(int i=2; i<=tot; i++) {
		res+=len[i]-len[fa[i]];
	}
	cout<<res<<endl;
}