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

 1 #include <bits/stdc++.h>
 2 #define il inline
 3 #define RG register
 4 #define ll long long
 5 #define M (4000010)
 6 #define N (100010)
 7 
 8 using namespace std;
 9 
10 struct edge{ int nt,to; }g[N<<1];
11 
12 int ch[M][10],fa[M],l[M],q[M],vis[M],la,tot;
13 int d[N],s[N],nxt[N],head[N],S,n,num;
14 ll ans;
15 
16 il int gi(){
17   RG int x=0,q=1; RG char ch=getchar();
18   while ((ch<‘0‘ || ch>‘9‘) && ch!=‘-‘) ch=getchar();
19   if (ch==‘-‘) q=-1,ch=getchar();
20   while (ch>=‘0‘ && ch<=‘9‘) x=x*10+ch-48,ch=getchar();
21   return q*x;
22 }
23 
24 il void insert(RG int from,RG int to){
25   g[++num]=(edge){head[from],to},head[from]=num; return;
26 }
27 
28 il int add(RG int p,RG int c){
29   RG int np=ch[p][c];
30   if (np && l[np]==l[p]+1) return np;
31   l[np=++tot]=l[p]+1;
32   for (;p && !ch[p][c];p=fa[p]) ch[p][c]=np;
33   if (!p){ fa[np]=1; return np; } RG int q=ch[p][c];
34   if (l[q]==l[p]+1) fa[np]=q; else{
35     RG int nq=++tot; l[nq]=l[p]+1;
36     fa[nq]=fa[q],fa[q]=fa[np]=nq;
37     memcpy(ch[nq],ch[q],sizeof(ch[q]));
38     for (;ch[p][c]==q;p=fa[p]) ch[p][c]=nq;
39   }
40   return np;
41 }
42 
43 il void dfs(RG int x,RG int p){
44   nxt[x]=add(nxt[p],s[x]);
45   for (RG int i=head[x];i;i=g[i].nt)
46     if (g[i].to!=p) dfs(g[i].to,x);
47   return;
48 }
49 
50 il void bfs(){
51   RG int h=0,t=1; q[t]=1,vis[1]=1;
52   while (h<t){
53     RG int x=q[++h],v;
54     for (RG int i=0;i<S;++i){
55       v=ch[x][i]; if (!v || vis[v]) continue;
56       ans+=l[v]-l[fa[v]],vis[v]=1,q[++t]=v;
57     }
58   }
59   return;
60 }
61 
62 int main(){
63 #ifndef ONLINE_JUDGE
64   freopen("substring.in","r",stdin);
65   freopen("substring.out","w",stdout);
66 #endif
67   n=gi(),S=gi(),nxt[0]=la=++tot;
68   for (RG int i=1;i<=n;++i) s[i]=gi();
69   for (RG int i=1,u,v;i<n;++i)
70     u=gi(),v=gi(),++d[u],++d[v],insert(u,v),insert(v,u);
71   for (RG int i=1;i<=n;++i) if (d[i]==1) dfs(i,0);
72   bfs(),cout<<ans; return 0;
73 }