代码

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