%%%陈老师%%%
简单来讲就是Trie上的后缀自动机,因为叶子节点总数不超过20,说以我们从每一个叶子节点出发做DFS,一边走一边把扫到的点甩进后缀自动机,这样后缀自动机里面的点不会超过4000000个,统计出本质不同的即可。
个人感觉Trie上写后缀自动机上的写法十分像树上的可持久化线段树的写法(SMG。。。先插入父亲再插入儿子
/**************************************************************
Problem: 3926
User: RicardoWang
Language: C++
Result: Accepted
Time:2368 ms
Memory:225364 kb
****************************************************************/
#include<cstdlib>
#include<cstdio>
#include<iostream>
#include<cstring>
#include<cmath>
#include<algorithm>
#include<queue>
#include<vector>
using namespace std;
#define WJMZBMR_OTZ true
#define maxa 4000005
#define maxn 100005
void _read(int &x)
{
x=0; char ch=getchar(); bool flag=false;
while(ch<'0'||ch>'9'){if(ch=='-')flag=false;ch=getchar();}
while(ch>='0'&&ch<='9'){x=x*10+ch-'0'; ch=getchar();}if(flag)x=-x;
return ;
}
struct edge
{
int to,next;
}e[maxn*2];
int edge_ct,head[maxn],col[maxn],n,c,du[maxn];
void add(int x,int y)
{
e[++edge_ct]=(edge){y,head[x]}; head[x]=edge_ct;
e[++edge_ct]=(edge){x,head[y]}; head[y]=edge_ct; return ;
}
long long ans;
struct Sam
{
int cnt,chi[maxa][12],len[maxa],fa[maxa];
Sam()
{
cnt=1;fa[1]=0;len[1]=0;
}
int extend(int last,int x)
{
int np=++cnt,p=last;len[np]=len[p]+1;
while(!chi[p][x]&&p)chi[p][x]=np,p=fa[p];
if(!p)fa[np]=1;
else
{
int q=chi[p][x];
if(len[q]==len[p]+1)fa[np]=q;
else
{
int nq=++cnt; len[nq]=len[p]+1;
memcpy(chi[nq],chi[q],sizeof(chi[nq]));
fa[nq]=fa[q];
fa[q]=fa[np]=nq;
while(chi[p][x]==q)chi[p][x]=nq,p=fa[p];
}
}
return np;
}
void DFS(int now,int fa,int last)
{
int t=extend(last,col[now]);
for(int id=head[now];id;id=e[id].next)
{
if(e[id].to==fa)continue;
DFS(e[id].to,now,t);
}
return ;
}
void solve()
{
for(int i=1;i<=cnt;i++)
{
ans+=len[i]-len[fa[i]];
}
return ;
}
}SAM;
void Init()
{
_read(n); _read(c);
int x,y;
for(int i=1;i<=n;i++)_read(col[i]);
memset(du,0,sizeof(du));
for(int i=1;i<n;i++)
{
_read(x); _read(y); du[x]++; du[y]++;
add(x,y);
}
}
int main()
{
//freopen("in.txt","r",stdin);
Init();
for(int i=1;i<=n;i++)if(du[i]==1)SAM.DFS(i,0,1);
SAM.solve();
cout<<ans<<endl;
return 0;
}