[백준] 15458 Barn Painting

문제

Farmer John has a large farm with �$N$ barns (1≤�≤105$1 \le N \le 10^5$), 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 �$N$ 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 �$N$ and �$K$ (0≤�≤�$0 \le K \le N$), respectively the number of barns on the farm and the number of barns that have already been painted.

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

The next �$K$ lines each contain two integers �$b$ and �$c$ (1≤�≤�$1 \le b \le N$, 1≤�≤3$1 \le c \le 3$) indicating that barn �$b$ is painted with color �$c$.

출력

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

 

 

 

N개의 Barn 이 n-1개의 edge로 연결되어있는 트리 형태입니다.

연결되어있는 barn은 색이 서로 달라야 하며, K개의 barn이 이미 색이 정해져있습니다.

트리에서 dp문제입니다.

한 정점의 색깔이 1일때, 경우의 수는 자식노드 2,3번 경우의수의 합의 곱입니다.

3차원 배열로 메모제이션을 사용합니다.

 

#include<bits/stdc++.h>
#define FAST ios_base::sync_with_stdio(false),cin.tie(NULL);
#define mset(v) memset(v,0,sizeof(v));
#define rep(i,a) for(int i=0;i<a;++i)
#define REP(i,a) for(int i=1;i<=a;++i)

using namespace std;

typedef long long ll;
typedef pair<int, int> pi;
typedef tuple<int, int, int>ti;
typedef vector<int> vi;
typedef vector<vector<int>> vvi;
int dy[] = { -1,0,1,0 }, dx[] = { 0,1,0,-1 }, INF = 987654321;

ll N, K, dp[100001][4], k[100001];
vi adj[100001];
int dfs(int n, int c, int p = 0) {
	if (dp[n][c] != -1)return dp[n][c];
	if (k[n]) {
		c == k[n] ? dp[n][c] = 1 : dp[n][c] = 0;
	}
	else {
		dp[n][c] = 1;
	}
	for (int r : adj[n]) {
		if (r==p)continue;
		ll sum = 0;
		for (int i = 1; i <= 3; i++) {
			if (c != i) {
				sum += dfs(r, i, n);
				sum %= 1000000007;
			}
		}
		dp[n][c] *= sum;
		dp[n][c] %= 1000000007;
	}
	return dp[n][c] % 1000000007;
}
int main(){
	FAST;
	cin >> N >> K;
	ll a, b;
	if (N == 0) { cout << 0; return 0; }
	rep(i, N - 1) {
		cin >> a >> b;
		adj[a].push_back(b);
		adj[b].push_back(a);
	}
	while (K--) {
		cin >> a >> b;
		k[a] = b;
	}
	memset(dp, -1, sizeof(dp));
	cout << (dfs(1, 1) + dfs(1, 2) + dfs(1, 3)) % 1000000007;

}