luogu3346 诸神眷顾的幻想乡 (广义SAM)

 1 #include<bits/stdc++.h>
 2 #define pa pair<int,int>
 3 #define CLR(a,x) memset(a,x,sizeof(a))
 4 #define MP make_pair
 5 using namespace std;
 6 typedef long long ll;
 7 const int maxn=1e5+10,maxp=4e6+10;
 8 
 9 inline char gc(){
10     return getchar();
11     static const int maxs=1<<16;static char buf[maxs],*p1=buf,*p2=buf;
12     return p1==p2&&(p2=(p1=buf)+fread(buf,1,maxs,stdin),p1==p2)?EOF:*p1++;
13 }
14 inline ll rd(){
15     ll x=0;char c=gc();bool neg=0;
16     while(c<‘0‘||c>‘9‘){if(c==‘-‘) neg=1;c=gc();}
17     while(c>=‘0‘&&c<=‘9‘) x=(x<<1)+(x<<3)+c-‘0‘,c=gc();
18     return neg?(~x+1):x;
19 }
20 
21 int N,C,eg[maxn*2][2],egh[maxn],ect,dgr[maxn];
22 int col[maxn];
23 int tr[maxp][12],fa[maxp],len[maxp],pct=1,rt=1;
24 
25 inline void adeg(int a,int b){
26     eg[++ect][0]=b,eg[ect][1]=egh[a],egh[a]=ect;
27     dgr[a]++;
28 }
29 
30 inline int insert(int x,int o){
31     int p=++pct;
32     len[p]=len[o]+1;
33     for(;o&&!tr[o][x];o=fa[o]) tr[o][x]=p;
34     if(!o){fa[p]=rt;return p;}
35     int q=tr[o][x];
36     if(len[q]==len[o]+1){fa[p]=q;return p;}
37     int qq=++pct;
38     memcpy(tr[qq],tr[q],sizeof(tr[qq]));
39     len[qq]=len[o]+1;fa[qq]=fa[q],fa[q]=fa[p]=qq;
40     for(;o&&tr[o][x]==q;o=fa[o]) tr[o][x]=qq;
41     return p;
42 }
43 
44 inline void dfs(int x,int f,int p){
45     p=insert(col[x],p);
46     for(int i=egh[x];i;i=eg[i][1]){
47         int b=eg[i][0];if(b==f) continue;
48         dfs(b,x,p);    
49     }
50 }
51 
52 int main(){
53     //freopen("","r",stdin);
54     int i,j,k;
55     N=rd(),C=rd();
56     for(i=1;i<=N;i++) col[i]=rd();
57     for(i=1;i<N;i++){
58         int a=rd(),b=rd();
59         adeg(a,b);adeg(b,a);
60     }
61       for(i=1;i<=N;i++){
62           if(dgr[i]==1) dfs(i,0,rt);
63       }
64       ll ans=0;
65       for(i=2;i<=pct;i++) ans+=len[i]-len[fa[i]];
66     printf("%lld\n",ans);
67     return 0;
68 }