3926: [Zjoi2015]诸神眷顾的幻想乡

//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;
}