boj, 백준) 15681. 트리와쿼리 ( C/ C++)

1. 문제 링크

www.acmicpc.net/problem/15681

15681번: 트리와 쿼리

2. 문제 개요

간선에 가중치와 방향성이 없는 임의의 루트 있는 트리가 주어졌을 때, 아래의 쿼리에 답해보도록 하자.

 - 정점 U를 루트로 하는 서브트리에 속한 정점의 수를 출력한다.

 

3. 문제 힌트

정점의 개수가 100,000개. 쿼리의 개수도 100,000개.

쿼리를 입력 받을 때, O(n^2)의 시간으로 처리했다가는 제 시간 안에 풀 수 없다. 미리 구해두고 쿼리가 들어오면 즉각 답하는 방법은 없을까? 

4. 문제 풀이

이 문제는 결국 쿼리로 들어온 노드를 기준으로하는 트리에 노드가 몇 개 있는지 구하면 되는 문제이다.

 

이 부분은 DFS로 해결할 수 있다. 트리 정보를 입력받고, 쿼리를 입력받기 전, 루트를 시작으로 DFS를 하여, 모든 노드에 대해, 자신을 루트로 하는 트리의 정점의 개수를 구할 수 있다.

 

DFS를 통해 자식이 없는 리프노트 까지 도달한 뒤, 해당 리프노드를 기준으로 하는 트리에서는 정점이 당연히 1개일 것이고, DFS를 반환하면서(부모로 올라감) 자신의 자식을 서브트리의 루트로 하는 서브트리의 정점의 수 + 1(자기 자신) 이런 방식으로 거슬러 올라간다.

 

이 부분은 말로 설명하기 보다는 코드를 보는 게 더 직관적이며, 아! 하고 이해하면 정말 쉽다.

처음엔.. 이게 무슨소리지? 라고 할 수 있다!!

 

이렇게 O(N)만큼의 시간을 소비해서 DFS로 탐색을 하고, 쿼리가 들어올 때 DFS를 하면서 저장했던 값을 통해 읽어주기만 하면 쿼리를 O(1)의 시간에 처리할 수 있다.

5. 코드

#include <cstdio>
#include <vector>
using namespace std;

vector<vector<int>> adj;
int n, r, q;
vector<int> nodenum;

int dfs(int parent, int cur)
{
	int subtreenum = 0;

	for (int next : adj[cur])
	{
		if (next == parent)
			continue;

		subtreenum += dfs(cur, next);
	}
	return nodenum[cur] = subtreenum + 1;
}

int main()
{
	scanf("%d %d %d", &n, &r, &q);
	adj.resize(n + 1);

	for (int i = 0; i < n - 1; ++i) {
		int u, v;
		scanf("%d %d", &u, &v);

		adj[u].push_back(v);
		adj[v].push_back(u);
	}
	nodenum.resize(n + 1);

	dfs(0, r);

	for (int j = 0; j < q; ++j) {
		int subroot;
		scanf("%d", &subroot);

		printf("%d\n", nodenum[subroot]);
	}

	return 0;
}

6. 결과