P4084

[USACO17DEC]Barn Painting G

题面翻译

题意：给定一颗 \(N\) 个节点组成的树，你要给每个点涂上三种颜色之一，其中有 \(K\) 个节点已染色，要求任意两相邻节点颜色不同，求合法染色方案数。

样例 #1

样例输入 #1

4 1
1 2
1 3
1 4
4 3

样例输出 #1

8

注意 f[u][1]=f[u][1]x(f[v][2]+f[v][3]) f[u][1]不止更新一次 所以要乘法原理 uの每个v独立！

#include<bits/stdc++.h>
using namespace std;
#define int long long
const int mod=1e9+7;
const int N=1e5+5;
int n,k;
struct Tree {
	int nxt,to;
} edge[N<<1];
int head[N],cnt,a[N],f[N][4];
void add(int u,int v) {
	cnt++;
	edge[cnt].to=v;
	edge[cnt].nxt=head[u];
	head[u]=cnt;
}
void dp(int u,int fat) { //f[u][col]
	if(!a[u])f[u][1]=f[u][2]=f[u][3]=1;
	if(a[u]==1)f[u][1]=1,f[u][2]=0,f[u][3]=0;
	if(a[u]==2)f[u][2]=1,f[u][1]=0,f[u][3]=0;
	if(a[u]==3)f[u][3]=1,f[u][1]=0,f[u][2]=0;
	for(int i=head[u]; i; i=edge[i].nxt) {
		int v=edge[i].to;
		if(v==fat)continue;
		dp(v,u);
		if(!a[u]) {
			f[u][1]=f[u][1]*(f[v][2]+f[v][3]);
			f[u][2]=f[u][2]*(f[v][1]+f[v][3]);
			f[u][3]=f[u][3]*(f[v][1]+f[v][2]);
			f[u][1]%=mod;
			f[u][2]%=mod;
			f[u][3]%=mod;
		} else {
			if(a[u]==1)f[u][1]=f[u][1]*(f[v][2]+f[v][3]),f[u][1]%=mod;
			if(a[u]==2)f[u][2]=f[u][2]*(f[v][1]+f[v][3]),f[u][2]%=mod;
			if(a[u]==3)f[u][3]=f[u][3]*(f[v][1]+f[v][2]),f[u][3]%=mod;
		}
	}
}
signed main() {
	ios::sync_with_stdio(false);
	cin>>n>>k;
	for(int i=1; i<n; i++) {
		int u,v;
		cin>>u>>v;
		add(u,v);
		add(v,u);
	}
	for(int i=1; i<=k; i++) {
		int b,c;
		cin>>b>>c;
		a[b]=c;
	}
	dp(1,0);
	cout<<(f[1][1]+f[1][2]+f[1][3])%mod<<"\n";
	return 0;
}