백준 15681번: 트리와 쿼리 (C++)

https://www.acmicpc.net/problem/15681

풀이방법

해당 문제는 아래 설명되어있는 힌트를 토대로 문제를 해결하였다. 문제를 해결한 sequence는 다음과 같다.

• 우선 입력 받은 r값을 루트로 하는 트리를 구현한다.
• 이때 탐색하는 현재 노드의 자식 노드를 child변수를 이용하여 저장해 둔다.
• countSubtreeNodes함수를 사용하여 r기준으로 만들어지는 트리를 바탕으로 각 노드별로 만들 수 있는 서브 트리의 개수를 구하여 treeSize 배열 변수에 저장한다.
• 최종적으로 입력 받은 UTree를 treeSize배열 인덱스로 사용하여 원하는 출력 값을 얻어낸다.

코드

#include <bits/stdc++.h>
using namespace std;


int n, r, q, u, v;
vector<int> adj[100005]; //해당 인덱스를 가지는 노드에 연결된 노드의 값
vector<int> child[100005]; // 해당 인덱스를 가지는 노드의 자식 노드들의 값
int treeSize[100005]; // 해당 인덱스를 루트로 가지는 노드의 자기자신 노드와 자식노드의 개수를 더한 값


void makeTree(int cur, int parent) // r기준으로 트리를 한번 만드는 함수
{
	for (auto nxt : adj[cur])
	{
		if (nxt == parent) continue;
		child[cur].push_back(nxt);
		makeTree(nxt, cur);
	}
}

void countSubtreeNodes(int cur) // r기준으로 만들어지는 트리를 바탕으로 각 노드별로 만들 수 있는 서브트리의 개수를 구하는 함수
{
	treeSize[cur] = 1;
	if (child[cur].empty()) return;
	for (auto nxt : child[cur])
	{
		countSubtreeNodes(nxt);
		treeSize[cur] += treeSize[nxt];
	}
}

int main() 
{
	ios::sync_with_stdio(0);
	cin.tie(0);
	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);
	}
	makeTree(r, -1);
	countSubtreeNodes(r);

	int UTree; 
	for (int i = 0; i < q; i++)
	{
		cin >> UTree;
		cout << treeSize[UTree] << '\n';
	}
}