代码

#include<cstdio>#include<cstdlib>#include<cstring>#include<iostream>using namespace std;typedef long long LL;const int N = 20 * 100010;int n, c, tot, len, last;int w[N], son[N][15], step[N], pre[N], first[N], cnt[N], id[N];struct node{int x, y, next;}a[2 * N];void ins(int x, int y){a[++len].x = x; a[len].y = y;a[len].next = first[x]; first[x] = len;}int add_node(int x){step[++tot] = x;return tot;}int extend(int p, int ch){int np = son[p][ch];if (np){if (step[np] == step[p] + 1){return np;}else{int temp = add_node(step[p] + 1);memcpy(son[temp], son[np], sizeof(son[np]));pre[temp] = pre[np]; pre[np] = temp;while (son[p][ch] == np)son[p][ch] = temp, p = pre[p]; return temp;}}//这是不同的部分np = add_node(step[p] + 1);//下面相同while (p && !son[p][ch]) son[p][ch] = np, p = pre[p];if (p == 0) pre[np] = 1;else{int q = son[p][ch];if (step[q] == step[p] + 1) pre[np] = q;else{int nq = add_node(step[p] + 1);memcpy(son[nq], son[q], sizeof(son[q]));pre[nq] = pre[q];pre[np] = pre[q] = nq;while (son[p][ch] == q) son[p][ch] = nq, p = pre[p];}}last = np;return np;}void dfs(int x, int fa, int now){int nt = extend(now, w[x]);// printf("%d\n",nt);for (int i = first[x]; i; i = a[i].next){int y = a[i].y;if (y != fa) dfs(y, x, nt);}}int main(){// freopen("a.in", "r", stdin);scanf("%d%d", &n, &c);for (int i = 1; i <= n; i++) scanf("%d", &w[i]);tot = 0; len = 0;memset(son, 0, sizeof(son));memset(pre, 0, sizeof(pre));memset(cnt, 0, sizeof(cnt));memset(first, 0, sizeof(first));step[++tot] = 0; last = 1;for (int i = 1; i<n; i++){int x, y;scanf("%d%d", &x, &y);ins(x, y); ins(y, x);cnt[x]++; cnt[y]++;}// for(int i=1;i<=len;i++) printf("%d -- > %d\n",a[i].x,a[i].y);for (int i = 1; i <= n; i++){if (cnt[i] == 1) dfs(i, 0, 1);}// for(int i=1;i<=tot;i++) printf("%d ",id[i]);printf("\n");LL ans = 0;for (int i = 1; i <= tot; i++) ans += (LL)(step[i] - step[pre[i]]);printf("%lld\n", ans);return 0;}