[BOJ] 15681번 트리와 쿼리 (C++)

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

15681번: 트리와 쿼리

 

#include <iostream>
#include <vector>
#define MAX 100001
using namespace std;

int dp[MAX];
vector<int> vec[MAX];
bool visited[MAX];

void dfs(int node, int parent){
	visited[node] = true;
	for(int i = 0; i < vec[node].size(); i++){
		int next = vec[node][i];
		if(!visited[next])
			dfs(next, node);
	}
	if(parent != -1){
		dp[parent] += dp[node];
	}
}

int main() {
	ios::sync_with_stdio(false);
	cin.tie(NULL);
	
	int n, r, q;
	cin >> n >> r >> q;

	for(int i = 0; i < n + 1; i++)
		dp[i] = 1;
		
	for(int i = 0; i < n - 1; i++){
		int u, v;
		cin >> u >> v;
		vec[u].push_back(v);
		vec[v].push_back(u);
	}
	
	dfs(r, -1);
	
	for(int i = 0; i < q; i++){
		int u;
		cin >> u;
		cout << dp[u] << endl;
	}
	return 0;
}

 

DFS + DP를 이용한 메모아이징을 사용한다. 

 

리프 노드는 자식이 없으므로 서브 트리의 정점 개수는 자기 자신 1개이다. 따라서 dp 배열을 1로 초기화한다.

루트 노드부터 dfs를 실행하고, 서브 트리 탐색이 끝났다면 부모에게 자신의 서브 트리 개수를 더해준다.

(부모가 -1인 경우는 루트 노드이므로 더하지 않는다.)

 

아래 예제에서 파란색으로 적은 숫자는 정점에서의 서브 트리 정점 개수이다.