[백준/BOJ] 백준 15458번 : Barn Painting

www.acmicpc.net/problem/15458

15458번: Barn Painting

루트가 1인 트리를 만들고, 트리 DP(트리에서 다이나믹 프로그래밍)를 이용해 문제를 해결했다. 주의해야 될 점은 there에 here에 칠하려는 색깔과 같은색이 칠해져 있는지를 고려해야 된다는 점이다.

 

코드

#include <iostream>
#include <algorithm>
#include <vector>
using namespace std;

int n, k;
vector<int> adj[100001];
vector<int> color_info(100001, 0);
vector<int> maked(100001, 0);
long long cache[100001][4];

void Pre()
{
	for (int i = 0; i < 100001; i++)
		for (int j = 0; j < 4; j++)
			cache[i][j] = -1;
}

struct node {

	int number;
	int color;
	vector<node*> children;

};

node* Make_tree(int root)
{
	node* ret = new node();

	ret->number = root;
	ret->color = color_info[root];
	maked[root] = 1;

	for (int i = 0; i < adj[root].size(); i++)
	{
		int there = adj[root][i];

		if (maked[there] == 0)
			ret->children.push_back(Make_tree(there));
	}

	return ret;
}

//here노드를 paint색으로 칠할때의 방법의 수
long long Solve(node* here, int paint)
{
	long long& ret = cache[here->number][paint];

	if (ret != -1)
		return ret;

	//색깔이 칠해져 있는데 칠하려는 색깔과 다를때
	if (here->color != 0 && here->color != paint)
	{
		ret = 0;
		return ret;
	}

	//leaf노드일때
	if (here->children.size() == 0)
	{
		ret = 1;
		return ret;
	}

	ret = 1;

	if (paint == 1)
	{
		for (int i = 0; i < here->children.size(); i++)
		{
			node* there = here->children[i];

			//there에 여기 칠하려는 색깔과 같은색이 있을때
			if (there->color == paint)
			{
				ret = 0;
				return ret;
			}

			ret = (((ret * (Solve(there, 2) + Solve(there, 3))) + 1000000007) % 1000000007);
		}
	}

	else if (paint == 2)
	{
		for (int i = 0; i < here->children.size(); i++)
		{
			node* there = here->children[i];

			if (there->color == paint)
			{
				ret = 0;
				return ret;
			}

			ret = (((ret * (Solve(there, 1) + Solve(there, 3))) + 1000000007) % 1000000007);
		}
	}

	else if (paint == 3)
	{
		for (int i = 0; i < here->children.size(); i++)
		{
			node* there = here->children[i];

			if (there->color == paint)
			{
				ret = 0;
				return ret;
			}

			ret = (((ret * (Solve(there, 1) + Solve(there, 2))) + 1000000007) % 1000000007);
		}
	}
	ret %= 1000000007;

	return ret;
}

int main()
{
	cin.tie(NULL);
	ios_base::sync_with_stdio(false);

	Pre();

	cin >> n >> k;

	for (int i = 0; i < n - 1; i++)
	{
		int u, v;

		cin >> u >> v;

		adj[u].push_back(v);
		adj[v].push_back(u);
	}

	for (int i = 0; i < k; i++)
	{
		int b, c;

		cin >> b >> c;

		color_info[b] = c;
	}

	//루트노드가 1인 트리를 만든다
	node* tree = Make_tree(1);

	long long result = 0;
	for (int i = 1; i <= 3; i++) //루트 노드를 1~3으로 색칠했을때 방법수를 모두 더한다
	{
		result = (((result + Solve(tree, i)) + 1000000007) % 1000000007);
	}

	cout << result;

	return 0;
}