先是广义后缀自动机的吧！
1 #include<cstdio> 2 #include<iostream> 3 #include<algorithm> 4 #include<cstring> 5 #define rep(i,j,k) for(register int i = j; i <= k; i++) 6 #define dow(i,j,k) for(register int i = j; i >= k; i--) 7 #define ll long long 8 using namespace std; 9 10 inline int read() { 11 int s = 0, t = 1; char c = getchar(); 12 while( !isdigit(c) ) { if( c == '-' ) t = -1; c = getchar(); } 13 while( isdigit(c) ) s = s * 10 + c - 48, c = getchar(); 14 return s * t; 15 } 16 17 const int N = 1e5+5, M = 4000200; 18 int v, nv, p, du[N], fa[M], dep[M], son[M][10], tot = 1; 19 struct edge{ int to; edge*nxt; } e[N<<1], *pt = e, *head[N]; 20 inline void add(int x,int y) { 21 du[x]++, du[y]++; 22 pt->to = y, pt->nxt = head[x], head[x] = pt++; 23 pt->to = x, pt->nxt = head[y], head[y] = pt++; 24 } 25 26 inline int extend(int u,int t) { 27 if( v = son[u][t] ) { 28 if( dep[v] == dep[u] + 1 ) return v; 29 nv = ++tot, dep[nv] = dep[u] + 1; 30 fa[nv] = fa[v], fa[v] = nv; 31 memcpy(son[nv],son[v],sizeof(son[v])); 32 while( son[u][t] == v ) son[u][t] = nv, u = fa[u]; 33 return nv; 34 } 35 p = ++tot, dep[p] = dep[u] + 1; 36 while( u && !son[u][t] ) son[u][t] = p, u = fa[u]; 37 if( !u ) fa[p] = 1; 38 else { 39 if( dep[v = son[u][t]] == dep[u] + 1 ) fa[p] = v; 40 else { 41 nv = ++tot, dep[nv] = dep[u] + 1; 42 fa[nv] = fa[v], fa[v] = fa[p] = nv; 43 memcpy(son[nv],son[v],sizeof(son[v])); 44 while( son[u][t] == v ) son[u][t] = nv, u = fa[u]; 45 } 46 } 47 return p; 48 } 49 50 #define ez(i,j) for(edge*i = head[j]; i; i=i->nxt) 51 #define to i->to 52 int c[N]; 53 inline void dfs(int x,int f,int now) { 54 now = extend(now,c[x]); 55 ez(i,x) if( to != f ) dfs(to,x,now); 56 } 57 58 int main() { 59 // freopen("in.txt","r",stdin); 60 int n = read(), m = read(); ll ans = 0; 61 rep(i,1,n) c[i] = read(); 62 rep(i,1,n-1) add(read(),read()); 63 rep(i,1,n) if( du[i] == 1 ) dfs(i,0,1); 64 rep(i,1,tot) ans += dep[i] - dep[fa[i]]; 65 cout<<ans<<endl; 66 return 0; 67 }