백준 / 15681 트리와 쿼리 C++

www.acmicpc.net/problem/15681

15681번: 트리와 쿼리

주어진 트리에서 입력 노드의 서브트리의 노드 개수를 구하는 문제입니다. 

c++로 풀다보니 입출력을 cin, cout으로 그냥 했는데 그냥 cin, cout으로 처리하면 입출력이 느려 시간초과가 났습니다.

(scanf, printf 이용하거나, ios::sync_with_stdio(false); cin.tie(NULL); cout.tie(NULL); 을 이용하면 된다.)

 

문제는 간단합니다.

dp 배열을 해당 노드의 서브 트리의 노드 개수라 합시다.

트리로 주어진 애들에 대해서 dfs를 루트 노드부터 수행합니다.

자식 노드의 dp값을 부모 노드에 더해주면 됩니다. (자신도 노드 개수에 포함해야 하므로 dp를 1로 초기화)

 

코드는 다음과 같습니다.

#include<iostream>
#include<algorithm>
#include<vector>
#include<cstring>

using namespace std;

int n, r, q;
vector<int>tree[100001];
int dp[100001]; // i번째 노드의 서브트리의 노드 수
bool check[100001];

void dfs(int node)
{
	check[node] = true;
	for (int i = 0; i < tree[node].size(); i++) {
		int next_node = tree[node][i];
		if (!check[next_node]) {
			dfs(next_node);
			dp[node] += dp[next_node];
		}
	}
}

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

	cin >> n >> r >> q;
	for (int i = 1; i < n; i++)
	{
		int a, b;
		cin >> a >> b;
		tree[a].push_back(b);
		tree[b].push_back(a);
		dp[a] = 1;
		dp[b] = 1;
	}

	dfs(r);

	while (q--)
	{
		int node;
		cin >> node;
		cout << dp[node] << '\n';
	}
}