自动搬运

查看原文

来自洛谷，原作者为

	teylnol_evteyl 我是狸猫盘的猫 | 哇，Teylnol Evteyl！ | 很慢地摇头

搬运于2025-08-24 22:49:56，当前版本为作者最后更新于2023-08-21 11:25:01，作者可能在搬运后再次修改，您可在原文处查看最新版

自动搬运只会搬运当前题目点赞数最高的题解，您可前往洛谷题解查看更多

以下是正文

---

设 $c_u$ 表示节点 $u$ 的颜色，$f_u$ 表示只考虑原树中 $u$ 的子树中的点、选择点 $u$ 的方案数。对于儿子 $v$，可以选择同色儿子节点，也可以跳过这个儿子节点，选择 $v$ 的与 $u$ 同色的儿子节点 $w$，故状态转移方程为：

$$f_u=\prod[c_u=c_v]f_v+\left(\prod[c_u=c_w]f_w+1\right) $$

其中 $v$ 是 $u$ 的儿子，$w$ 是 $v$ 的儿子。

统计答案时，除了考虑选择的点的最近公共祖先被选择的情况（即 $\sum f_u$），还有最近公共祖先没有被选择的情况，也就是说一个节点 $u$ 的几个儿子节点颜色相同，则可以分别在它们的子树中选择，而不选 $u$。设 $g_{u,c}$ 表示考虑 $u$ 的子树，不选择 $u$，选择至少两个颜色为 $c$ 的儿子节点的方案数，因为要分别减去选择 $0$ 或 $1$ 个儿子的情况，所以：

$$g_{u,c}=\left(\prod[c_v=c]f_v+1\right)-\left(\sum[c_v=c]f_v\right)-1 $$

其中 $v$ 是 $u$ 的儿子。

#include <bits/stdc++.h>

using namespace std;

typedef long long LL;

const int N = 1e5 + 5, P = 1e9 + 7;

int n, c[N];
int la[N], ne[N * 2], en[N * 2], idx;
LL f[N], g[N], res;

void add(int a, int b)
{
	ne[ ++ idx] = la[a];
	la[a] = idx;
	en[idx] = b;
}
void dp(int u, int fa)
{
	f[u] = 1;
	for(int i = la[u]; i; i = ne[i])
	{
		int v = en[i];
		if(v == fa) continue ;
		dp(v, u);
		
		LL t = 1;
		for(int j = la[v]; j; j = ne[j])
		{
			int w = en[j];
			if(w == u) continue ;
			if(c[w] == c[u]) t = t * (f[w] + 1) % P;
		}
		if(c[v] == c[u]) t = (t + f[v]) % P;
		f[u] = f[u] * t % P;
	}
	for(int i = la[u]; i; i = ne[i])
	{
		int v = en[i];
		if(v == fa) continue ;
		g[c[v]] = g[c[v]] * (f[v] + 1) % P;
	}
	for(int i = la[u]; i; i = ne[i])
	{
		int v = en[i];
		if(v == fa) continue ;
		res = (res + g[c[v]] - 1) % P;
		g[c[v]] = 1;
	}
}

int main()
{
	scanf("%d", &n);
	for(int i = 1; i <= n; i ++ ) scanf("%d", &c[i]);
	for(int i = 1; i < n; i ++ )
	{
		int a, b;
		scanf("%d%d", &a, &b);
		add(a, b), add(b, a);
	}
	add(0, 1);
	for(int i = 1; i <= n; i ++ ) g[i] = 1;
	dp(0, 0);
	printf("%lld\n", res);
	return 0;
}