Luogu-3346 [ZJOI2015]诸神眷顾的幻想乡
\(trie\)树建广义后缀自动机:
\(dfs\)遍历\(trie\)树,将树上的一个节点插入\(sam\)时,将他的\(fa\)在\(sam\)上所在的节点作为\(last\)
#include<map>
#include<queue>
#include<cmath>
#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
typedef long long ll;
const int maxn=2e5+100,maxm=4e6+100;
struct SAM{
int son[maxm][11],len[maxm],fa[maxm];
int last,tot,head[maxn],num,nex[maxn],v[maxn],c[maxn],du[maxn];
SAM(){last=tot=0,fa[0]=-1,num=1;}
void add(int x,int y){
v[++num]=y;
nex[num]=head[x];
head[x]=num;
v[++num]=x;
nex[num]=head[y];
head[y]=num;
du[x]++,du[y]++;
}
int insert(int x,int last){
int p=last,np=son[p][x];
if(np&&len[np]==len[p]+1) return np;
np=++tot;
len[np]=len[p]+1;
while(~p&&!son[p][x])
son[p][x]=np,p=fa[p];
if(p==-1)
fa[np]=0;
else{
int q=son[p][x];
if(len[q]==len[p]+1)
fa[np]=q;
else{
int nq=++tot;
memcpy(son[nq],son[q],sizeof(son[q]));
fa[nq]=fa[q];
len[nq]=len[p]+1;
fa[q]=fa[np]=nq;
while(~p&&son[p][x]==q)
son[p][x]=nq,p=fa[p];
}
}
return np;
}
void dfs(int x,int fa,int last){
last=insert(c[x],last);
for(int i=head[x];i;i=nex[i])
if(v[i]!=fa)
dfs(v[i],x,last);
}
void query(){
ll ans=0;
for(int i=1;i<=tot;i++)
ans+=1ll*(len[i]-len[fa[i]]);
printf("%lld\n",ans);
}
}sam;
int n,m,a,b;
int main(){
scanf("%d%d",&n,&m);
for(int i=1;i<=n;i++) scanf("%d",&sam.c[i]);
for(int i=1;i<n;i++)
scanf("%d%d",&a,&b),sam.add(a,b);
for(int i=1;i<=n;i++)
if(sam.du[i]==1)
sam.dfs(i,i,0);
sam.query();
return 0;
}