一个广义后缀自动机模板。
//Achen
#include<algorithm>
#include<iostream>
#include<cstring>
#include<cstdlib>
#include<vector>
#include<cstdio>
#include<queue>
#include<cmath>
const int N=4000007;
typedef long long LL;
using namespace std;
int n,c,col[N];
template<typename T>void read(T &x) {
char ch=getchar(); x=0; T f=1;
while(ch!=‘-‘&&(ch<‘0‘||ch>‘9‘)) ch=getchar();
if(ch==‘-‘) f=-1,ch=getchar();
for(;ch>=‘0‘&&ch<=‘9‘;ch=getchar()) x=x*10+ch-‘0‘; x*=f;
}
int ecnt,fir[N],nxt[N],to[N],in[N];
void add(int u,int v) {
nxt[++ecnt]=fir[u]; fir[u]=ecnt; to[ecnt]=v; in[v]++;
nxt[++ecnt]=fir[v]; fir[v]=ecnt; to[ecnt]=u; in[u]++;
}
int tot,rt,p,np,ch[N][11],fa[N],sz[N],l[N];
int insert(int c,int last) {
p=last; np=++tot;
l[np]=l[p]+1;
for(;p&&!ch[p][c];p=fa[p]) ch[p][c]=np;
if(!p) fa[np]=rt;
else {
int q=ch[p][c];
if(l[p]+1==l[q]) fa[np]=q;
else {
int nq=++tot; l[nq]=l[p]+1;
memcpy(ch[nq],ch[q],sizeof(ch[q]));
fa[nq]=fa[q]; fa[q]=fa[np]=nq;
for(;p&&ch[p][c]==q;p=fa[p]) ch[p][c]=nq;
}
}
sz[np]++; return np;
}
void dfs(int x,int fa,int pa) {
int last=insert(col[x],pa);
for(int i=fir[x];i;i=nxt[i]) if(to[i]!=fa) {
dfs(to[i],x,last);
}
}
int main() {
read(n); read(c); rt=++tot;
for(int i=1;i<=n;i++) read(col[i]);
for(int i=1;i<n;i++) {
int u,v;
read(u); read(v);
add(u,v);
}
for(int i=1;i<=n;i++) if(in[i]==1)
dfs(i,0,rt);
LL ans=0;
for(int i=2;i<=tot;i++) ans+=(LL)l[i]-l[fa[i]];
printf("%lld\n",ans);
return 0;
}