[bzoj3926]诸神眷顾的幻想乡

 1 #include<bits/stdc++.h>
 2 using namespace std;
 3 #define N 100005
 4 struct ji{
 5     int nex,to;
 6 }edge[N<<1];
 7 int E,V,n,x,y,head[N],a[N],sum[N*40],len[N*40],fa[N*40],ch[N*40][11];
 8 long long ans;
 9 char s[N];
10 void add1(int x,int y){
11     edge[E].nex=head[x];
12     edge[E].to=y;
13     head[x]=E++;
14 } 
15 int add2(int c,int last){
16     int p=last,np=last=++V;
17     len[np]=len[p]+1;
18     sum[np]=1;
19     for(;(p)&&(!ch[p][c]);p=fa[p])ch[p][c]=V;
20     if (!p)fa[np]=1;
21     else{
22         int q=ch[p][c];
23         if (len[q]==len[p]+1)fa[np]=q;
24         else{
25             int nq=++V;
26             len[nq]=len[p]+1;
27             memcpy(ch[nq],ch[q],sizeof(ch[q]));
28             fa[nq]=fa[q];
29             fa[q]=fa[np]=nq;
30             for(;(p)&&(ch[p][c]==q);p=fa[p])ch[p][c]=nq;
31         }
32     }
33     return last;
34 }
35 void dfs(int k,int fa,int l1){
36     int l2=add2(a[k],l1);
37     for(int i=head[k];i!=-1;i=edge[i].nex)
38         if (edge[i].to!=fa)dfs(edge[i].to,k,l2);
39 }
40 int main(){
41     scanf("%d%*d",&n);
42     V=1;
43     memset(head,-1,sizeof(head));
44     for(int i=1;i<=n;i++)scanf("%d",&a[i]);
45     for(int i=1;i<n;i++){
46         scanf("%d%d",&x,&y);
47         add1(x,y);
48         add1(y,x);
49     }
50     for(int i=1;i<=n;i++)
51         if (edge[head[i]].nex==-1)dfs(i,0,1);
52     for(int i=1;i<=V;i++)ans+=len[i]-len[fa[i]];
53     printf("%lld",ans);
54 }