bzoj3926: [Zjoi2015]诸神眷顾的幻想乡

 1 #include <iostream>
 2 #include <cstdio>
 3 #include <cstring>
 4 #include <cmath>
 5 #include <algorithm>
 6 #define maxn 100005
 7 #define maxl 200005
 8 #define maxm 4000005
 9 using namespace std;
10 
11 typedef long long ll;
12 int n,c,cnt,tot,root,last,fa[maxm],du[maxn],dist[maxm],son[maxl],prep[maxl],now[maxn],v[maxn];
13 ll ans;
14 void add(int x,int y){
15     cnt++,prep[cnt]=now[x],now[x]=cnt,son[cnt]=y;
16 }
17 struct Tsegment{
18     int son[maxm][10];
19     void prepare(){tot=last=root=1;memset(dist,0,sizeof(dist));}
20     int newnode(int x){dist[++tot]=x; return tot;}
21     int add(int x,int p){
22         int q=son[p][x];
23         if (q==0){
24             int np=newnode(dist[p]+1); last=np;
25             for (;p&&!son[p][x];p=fa[p]) son[p][x]=np;
26             if (p==0) fa[np]=root;
27             else{
28                 q=son[p][x];
29                 if (dist[p]+1==dist[q]) fa[np]=q;
30                 else{
31                     int nq=newnode(dist[p]+1);
32                     memcpy(son[nq],son[q],sizeof(son[q]));
33                     fa[nq]=fa[q],fa[q]=fa[np]=nq;
34                     for (;p&&son[p][x]==q;p=fa[p]) son[p][x]=nq;
35                 }
36             }
37         }else{
38             if (dist[p]+1==dist[q]) last=q;
39             else{
40                 int nq=newnode(dist[p]+1); last=nq;
41                 memcpy(son[nq],son[q],sizeof(son[q]));
42                 fa[nq]=fa[q],fa[q]=nq;
43                 for (;p&&son[p][x]==q;p=fa[p]) son[p][x]=nq;
44             }
45         }
46         return last;
47     }
48 }SAM;
49 void dfs(int x,int pa,int goal){
50     int p=SAM.add(v[x],goal);
51     for (int i=now[x],so=son[i];i;i=prep[i],so=son[i]){
52         if (so==pa) continue;
53         dfs(so,x,p);
54     }
55 }
56 int main(){
57     scanf("%d%d",&n,&c);
58     memset(du,0,sizeof(du));
59     cnt=tot=ans=0,memset(now,0,sizeof(now));
60     for (int i=1;i<=n;i++) scanf("%d",&v[i]);
61     for (int u,v,i=1;i<n;i++) scanf("%d%d",&u,&v),add(u,v),add(v,u),du[u]++,du[v]++;
62     SAM.prepare();
63     for (int i=1;i<=n;i++){
64         if (du[i]>1) continue;
65         dfs(i,0,root);
66     }
67     ans=0;
68     dist[root]=0;
69     for (int i=1;i<=tot;i++) ans+=(dist[i]-dist[fa[i]]);
70     printf("%lld\n",ans);
71     return 0;
72 }