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

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

15681번: 트리와 쿼리

 

 접근법

- 먼저 그래프가 주어지고 이것의 루트가 주어진다. 그리고 트리를 구성하게된다. 그래프가 주어지고서 이것을 트리로 바꾸어야하기 때문에, 이것을 문제의 그림으로 표현하면 다음과 같다. 

 

 

 

이 그래프에서 5를 루트라고 하였으므로 아래처럼 트리를 구성할 수 있다.

 

 

그리고 각 노드들의 서브트리의 노드 개수를 출력하면 된다. 이해하기는 쉽지만 어떻게 새어야할까. 먼저 어떤 것이 루트로 주어질지 모르기 때문에 모두 양방향으로 노드를 입력받는다. 그렇기 때문에 3을 구한다고 하였을 때 트리탐색을 하면 자신의 부모노드까지 방문하게 된다. 

 

그래서 모든 노드를 방문하며 서브트리노드의 개수를 새어 저장해두어야한다.

그러므로 이 문제는 트리에서 다이나믹 프로그래밍 이다.

 

dfs를 통해 끝까지 들어간 후 다시 빠져나오면서 dp를 갱신하는 것이 이 문제의 핵심이다. 그림으로 살펴보면 다음과 같다.

 

 

dfs의 특성상 한 방향으로 끝까지 들어간다. 위 사진을 보면 1, 2 번에서 처음 노드를 방문할 때 마다 dp[node] = 1 로 초기화한다. 왜냐면 마지막 리프노드라도 나 자신은 갖고 있기 때문이다. 

그리고 빨간색 화살표로 다시 dfs가 끝나고 다시 2번으로 돌아왔을 때, dp[2] = dp[2] + dp[4] 가 된다. 자기자신을 더하는 이유는 자식노드가 한개가 아닐수 있기 때문이다. 계속 더해줘야한다.

 

이것을 dfs 코드로 보면 다음과 같다.

 

void dfs(int Node) {
	dp[Node] = 1;
	visited[Node] = true;

	for (int i = 0; i < node[Node].size(); i++) {
		int nextNode = node[Node][i];
		if (visited[nextNode]) continue;
		visited[nextNode] = true;
		dfs(nextNode);
		dp[Node] += dp[nextNode];
	}
}

 

처음 1로 초기화하고 dfs가 끝났을 때 정보를 갱신한다. 이것이 트리에서 다이나믹 프로그래밍의 가장 기본적인 방법이다. 

 

정답 코드는 다음과 같다.

#include <iostream>
#include <vector>
#define MAX 100001
using namespace std;
int N, R, Q;

vector<int> node[MAX];
int dp[MAX];

bool visited[MAX] = { false, };
void dfs(int Node) {
	dp[Node] = 1;
	visited[Node] = true;

	for (int i = 0; i < node[Node].size(); i++) {
		int nextNode = node[Node][i];
		if (visited[nextNode]) continue;
		visited[nextNode] = true;
		dfs(nextNode);
		dp[Node] += dp[nextNode];
	}
}

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

	cin >> N >> R >> Q;

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

	dfs(R);

	for (int i = 0; i < Q; i++) {
		int tmp;
		cin >> tmp;
		cout << dp[tmp] << '\n';
	}
}