【洛谷P4084】Barn Painting【树形DP】

题目大意：

题目链接：https://www.luogu.org/problemnew/show/P4084
一棵$n$个节点的树上有$k$个点已被染色。求将这棵树染成三种颜色且相邻的节点颜色不同的方案数。

---

思路：

肯定是树形DP啊。设$f[i][1/2/3]$表示第$i$个节点颜色为$1/2/3$的时候以$i$为根的方案数。
那么考虑所有$i$的子树$son[i]$，由于这些子树肯定互不关联，所以任意一个子树对其它子树就没有影响。那么很明显的，第一个子树的所有情况可以和第二个子树的所有情况都匹配，所以就有$f[son[1]]\times f[son[2]]$种方案。然后加入第三颗子树，就有$f[son[1]]\times f[son[2]]\times f[son[3]]$种方法，以此类推。。。

所以就有：
$f[i][color1]=\prod^{子节点个数}_{j=1}(f[son[j]][color2]+f[con[j]][color3])$
别忘了要取模$10^9+7$。
最终答案为$(f[1][1]+f[1][2]+f[1][3])\%MOD$

---

代码：

#include <cstdio>
#include <cstring>
#include <iostream>
#define N 100100
#define MOD 1000000007ll
#define ll long long
using namespace std;

int n,m,tot,head[N],color[N];
ll f[N][4];
bool vis[N][4];  //记录是否访问过这个节点

struct edge
{
	int next,to;
}e[N*2];

void add(int from,int to)
{
	tot++;
	e[tot].to=to;
	e[tot].next=head[from];
	head[from]=tot;
}

ll dp(int x,int col,int fa)
{
	if (color[x]&&col!=color[x])
	 return 0;
	if (vis[x][col]) return f[x][col];  //剪枝
	vis[x][col]=true;
	ll ans=0;
	f[x][col]=1;
	for (int i=head[x];~i;i=e[i].next)
	 if (e[i].to!=fa)
	 {
	 	ans=0;
	 	for (int j=1;j<=3;j++)  //枚举每个颜色
	  	 if (j!=col)  //颜色要不同
	  	  ans=(ans+dp(e[i].to,j,x))%MOD;
	  	f[x][col]=(f[x][col]*ans)%MOD;
	 } 
	return f[x][col];
}

int main()
{
	memset(head,-1,sizeof(head));
	scanf("%d%d",&n,&m);
	int x,y;
	for (int i=1;i<n;i++)
	{
		scanf("%d%d",&x,&y);
		add(x,y);
		add(y,x);
	}
	for (int i=1;i<=m;i++)
	{
		scanf("%d%d",&x,&y);
		color[x]=y;
	}
	dp(1,1,-1);
	dp(1,2,-1);
	dp(1,3,-1);
	cout<<(f[1][1]+f[1][2]+f[1][3])%MOD;
	return 0;
}