[백준/BOJ] 백준 15681번 : 트리와 쿼리

www.acmicpc.net/problem/15681

15681번: 트리와 쿼리

트리에서 다이나믹 프로그래밍(트리DP)를 통해 문제를 해결했다. 루트의 번호가 r인 트리를 만들고, 트리의 정점의 수를 구하면서 cache에 각 노드가 루트인 서브 트리의 정점의 개수를 저장했다.

 

코드

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

vector<int> adj[100001];
vector<int> maked(100001, 0); //해당정점이 트리에 만들어졌는지 확인
vector<int> cache(100001, -1);

struct node
{
	int number;
	vector<node*> children;
};

//루트노드의 번호가 root인 트리를 만든다
node* treeMake(int root)
{
	node* ret = new node();
	ret->number = root;

	maked[root] = 1; //트리에 만든 정점 표시

	for (int i = 0; i < adj[root].size(); i++)
	{
		if (maked[adj[root][i]] == 0)
		{
			ret->children.push_back(treeMake(adj[root][i]));
		}
	}

	return ret;
}

//root트리에 속한 정점의 수를 구한다
int Solve(node* root)
{
	int& ret = cache[root->number];

	//이미 계산한적이 있을때
	if (ret != -1)
		return ret;

	//자식노드가 없을때
	if ((root->children).size() == 0)
	{
		ret = 1;
		return ret;
	}

	ret = 1; //현재의 정점

	for (int i = 0; i < (root->children).size(); i++)
	{
		//현재 노드의 자식노드가 루트인 서브트리의 정점의 개수를 더한다
		ret += Solve((root->children)[i]);
	}

	return ret;
}

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

	int n, r, q;
	int u, v;
	node* tree;

	cin >> n >> r >> q;

	for (int i = 0; i < n - 1; i++)
	{
		cin >> u >> v;
		adj[u].push_back(v);
		adj[v].push_back(u);
	}

	//루트의 번호가 r인 트리를 만든다
	tree = treeMake(r);

	//트리의 정점의 수를 구하면서 cache에 각 노드가 루트인 서브트리의 정점의 개수를 저장한다
	Solve(tree);

	for (int i = 0; i < q; i++)
	{
		cin >> u;
		cout << cache[u] << "\n";
	}

	return 0;
}