[洛谷]P4084 [USACO17DEC]Barn Painting (#树形dp)

题意翻译

题意：给定一颗N个节点组成的树，3种颜色，其中K个节点已染色，要求任意两相邻节点颜色不同，求合法染色方案数。 翻译贡献者：Il_ItzABC_lI

题目描述

Farmer John has a large farm with NN barns (1 \le N \le 10^51≤N≤105), some of which are already painted and some not yet painted. Farmer John wants to paint these remaining barns so that all the barns are painted, but he only has three paint colors available. Moreover, his prize cow Bessie becomes confused if two barns that are directly reachable from one another are the same color, so he wants to make sure this situation does not happen.

It is guaranteed that the connections between the NN barns do not form any 'cycles'. That is, between any two barns, there is at most one sequence of connections that will lead from one to the other.

How many ways can Farmer John paint the remaining yet-uncolored barns?

输入格式

The first line contains two integers NN and KK (0 \le K \le N0≤K≤N), respectively the number of barns on the farm and the number of barns that have already been painted.

The next N-1N−1 lines each contain two integers xx and yy (1 \le x, y \le N, x \neq y1≤x,y≤N,x≠y) describing a path directly connecting barns xx and yy.

The next KK lines each contain two integers bb and cc (1 \le b \le N1≤b≤N, 1 \le c \le 31≤c≤3) indicating that barn bbis painted with color cc.

输出格式

Compute the number of valid ways to paint the remaining barns, modulo 10^9 + 7109+7, such that no two barns which are directly connected are the same color.

输入输出样例

输入 #1复制

4 1
1 2
1 3
1 4
4 3

输出 #1复制

8

---

思路

树形dp好题。和没有丧尸的舞会那题有点像。

令dp[i][j]为在节点i染第j种颜色的方案数，把1号节点看成根。

由于相邻的点颜色不能相同，所以转移的时候只能由颜色不同的转移过来。由乘法原理，当前节点i的方案数应为其子树方案数的乘积。

dp[i][1]=dp[i][1]×(dp[son][2]+dp[son][3])

dp[i][2]=dp[i][2]×(dp[son][1]+dp[son][3])

dp[i][3]=dp[i][3]×(dp[son][1]+dp[son][2])

对于一个最初没有涂色的点i，刚遍历到的时候设置点i的初始值dp[i][1]=dp[i][2]=dp[i][3]=1。如果点i已经被上色了，那么只有被上的颜色的数组才能为1。

最后答案就是dp[root][1]+dp[root][2]+dp[root][3]了。

#include <stdio.h>
#include <iostream>
#define ll long long int 
#define mod 1000000007
#define maxn 100001
using namespace std;
ll n,m,s,cnt,head[maxn],color[maxn];
ll dp[maxn][4];
bool vis[maxn];
struct node
{
	ll nxt,to;
}e[maxn<<1];
inline void add(int u,int v)//前向星存图 
{
	e[++cnt].to=v;
	e[cnt].nxt=head[u];
	head[u]=cnt;
}
void dfs(int i)
{
	vis[i]=1;
	if(color[i])//如果该节点i已经上过色 
	{
		dp[i][color[i]]=1;//该种颜色都有，其他2种方案数为0 
	}
	else//如果没上过色 
	{
		dp[i][1]=1;//其他颜色都有可能被上 
		dp[i][2]=1;
		dp[i][3]=1;
	}
	for(register ll j=head[i];j;j=e[j].nxt)//找节点i的子节点 
	{
		int v(e[j].to);
		if(vis[v]==0)
		{
			dfs(v);
			dp[i][1]=dp[i][1]*((dp[v][2]+dp[v][3])%mod)%mod;
			dp[i][2]=dp[i][2]*((dp[v][1]+dp[v][3])%mod)%mod;
			dp[i][3]=dp[i][3]*((dp[v][1]+dp[v][2])%mod)%mod;
		}
	}
}
signed main()
{
	ios::sync_with_stdio(false);
	cin.tie(0);
	cout.tie(0);
	register ll i,j,u,v;
	cin>>n>>m;
	for(i=1;i<=n-1;i++)
	{
		cin>>u>>v;
		add(u,v);//无向 
		add(v,u);
	}
	for(i=1;i<=m;i++)
	{
		cin>>u>>v;
		color[u]=v;//u为节点，v为颜色种类 
	}
	dfs(1);//令节点1为根节点 
	cout<<(dp[1][1]+dp[1][2]+dp[1][3])%mod<<endl;
	return 0;
}