BZOJ3926: [Zjoi20150]诸神眷顾的幻想乡

#include<cstdio>
#include<iostream>
#include<cstring>
#include<cstdlib>
#include<algorithm>
using namespace std;

struct Node
{
int len;
Node *f,*last,*ch[10];
Node(){for(len=0;len<=9;len++)ch[len]=NULL;last=NULL;len=0;}
}*root,*last;

inline void add(int data)
{
Node *tp,*ne=new Node;
ne->len=last->len+1;
ne->f=last;
for(tp=last;tp&&!tp->ch[data];tp=tp->last)tp->ch[data]=ne;
if(!tp)
ne->last=root,last=ne;
else if(tp->len==tp->ch[data]->len-1)
ne->last=tp->ch[data],last=ne;
else
{
Node *a=new Node,*b=tp->ch[data];
*a=*b;
ne->last=b->last=a;
a->len=tp->len+1;
last=ne;
a->f=tp;
for(tp;tp&&tp->ch[data]==b;tp=tp->last)tp->ch[data]=a;
}
}
struct Chain
{
Chain*next;
int u;
Chain(){next=NULL;}
}*Head[1000001];

inline void addside(int a,int b)
{
Chain *tp=new Chain;
tp->u=b,tp->next=Head[a],Head[a]=tp;
}
int color[100001];
void DFS(int u,int fa)
{
for(Chain *tp=Head[u];tp;tp=tp->next)
if(fa^tp->u)
add(color[tp->u]),DFS(tp->u,u);
last=last->f;
}

int n,C;
char c;
inline void read(int &a)
{
a=0;do c=getchar();while(c<'0'||c>'9');
while(c<='9'&&c>='0')a=(a<<3)+(a<<1)+c-'0',c=getchar();
}
int Sort[1000001];
int H[1000001];
inline bool cmp(int x,int y)
{return H[x]<H[y];}
long long ans;
void Solve(Node *tp)
{
if(tp->last)
ans+=tp->len-tp->last->len;
for(int i=0;i<=9;i++)
if(tp->ch[i]&&tp->ch[i]->len==tp->len+1)
Solve(tp->ch[i]);
}
int main()
{
read(n),read(C);
int i,j,k;
for(i=1;i<=n;i++)
Sort[i]=i,read(color[i]);
root=last=new Node;
for(k=1;k<n;k++)
read(i),read(j),addside(i,j),addside(j,i),H[i]++,H[j]++;
sort(Sort+1,Sort+1+n,cmp);
for(i=1;i<=n;i++)
if(H[Sort[i]]==1)
{
last=root;
add(color[Sort[i]]);
DFS(Sort[i],0);
}
else break;
Solve(root);
printf("%lld\n",ans);
return 0;
}