洛谷_[P4084]Barn Painting G题解

题目链接
这题我们可以定义一个二维的 dp 数组，
在 dp[i][j]中：
i 表示对于节点 i，
j 有 1,2,3 三种状态，
表示当点 i 选择被染成颜色 j 时，以 i 为根的这颗子树有多少种染色方法。
那么根据乘法原理，节点 i 的方案数肯定等于 i 的每个儿子的方案数量之积。
这道题整个挺简单的，剩下细节的部分直接看下面的代码吧，（里面有详细注释）：

#include <bits/stdc++.h>
using namespace std;
#define int long long
#define mod 1000000007
int cnt = 0,head[100010],dp[100010][4],col[100010];
// dp 数组的意思就是我前面说的，col 数组用于记录该节点是不是提前被染成了什么颜色

struct Node
{
	int to,nex;
}e[200010];

void add(int x,int y)
{
	cnt++;
	e[cnt].to = y;
	e[cnt].nex = head[x];
	head[x] = cnt;
}

void dfs(int x,int fa)
{
	dp[x][1] = dp[x][2] = dp[x][3] = 1; // 因为后续要乘，所以默认初值是 1
	if(col[x] != 0) // 如果 x 本身有颜色
	{
		for(int i = 1; i <= 3; i++)
		{
			if(i != col[x])
				dp[x][i] = 0; // 那么这个节点不可能染成其他颜色
		}
	}
	bool flag = false; // 用于判断 x 是不是叶节点
	bool vis[4];
	vis[1] = vis[2] = vis[3] = false; // vis[i] 用于记录 x 的子树中有没有初始就染成颜色 i 的
	for(int i = head[x]; i; i = e[i].nex) // 前向星遍历
	{
		int y = e[i].to;
		if(y == fa) // 我们不能通过 x 又回到它的父节点去了
			continue;
		flag = true; // x 不是叶子节点
		if(col[y] != 0) // 如果 y 自带颜色
			vis[col[y]] = true; // 标记这种颜色
		dfs(y,x); // 进入这个儿子的 dfs
		if(col[x] != 0)
		{
			int tmp = 0;
			for(int j = 1; j <= 3; j++)
			{
				if(j == col[x]) // 如果 x 本身有颜色，则不能统计儿子中和它染成一样的颜色的方案
					continue;
				tmp += dp[y][j];
			}
			dp[x][col[x]] *= tmp; // 乘法原理
			dp[x][col[x]] %= mod;
		}
		else // 如果 x 本身没颜色，就没限制
		{
			// x 的子树不能染和它一样的颜色
			dp[x][1] *= dp[y][2] + dp[y][3];
			dp[x][1] %= mod;
			dp[x][2] *= dp[y][1] + dp[y][3];
			dp[x][2] %= mod;
			dp[x][3] *= dp[y][1] + dp[y][2];
			dp[x][3] %= mod;
		}
	}
	dp[x][1] %= mod;
	dp[x][2] %= mod;
	dp[x][3] %= mod;
	for(int i = 1; i <= 3; i++)
	{
		if(vis[i]) // 如果儿子中有这种颜色，则 x 不能染成这个颜色
			dp[x][i] = 0;
	}
}

signed main()
{
	int n,k;
	cin >> n >> k;
	for(int i = 1; i < n; i++)
	{
		int x,y;
		cin >> x >> y;
		add(x,y);
		add(y,x);
	}
	for(int i = 1; i <= k; i++)
	{
		int b,c;
		cin >> b >> c;
		col[b] = c;
	}
	dfs(1,-1); // 随便从哪个点开始，只不过我选的是 1 号点
	cout << (dp[1][1] + dp[1][2] + dp[1][3]) % mod << "\n"; // 最终答案是根节点染成颜色 1、2、3的方案数之和
	return 0;
}

记录
感谢阅读！求支持