「Daily OI Round 1」block 题解

标签：en OI la int 题解 ne 儿子 节点 block

设 \(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;
}

标签：en,OI,la,int,题解,ne,儿子,节点,block
From： https://www.cnblogs.com/recollect-the-past/p/17981305