C++ 백준 15681 트리와 쿼리

백준 15681

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

15681번: 트리와 쿼리

풀이 전 나의 생각

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

 

조건

- (2 ≤ N ≤ 10^5, 1 ≤ R ≤ N, 1 ≤ Q ≤ 10^5)

- (1 ≤ U, V ≤ N, U ≠ V)

- (1 ≤ U ≤ N)

 

과정

입력에 관한 조건들만 존재하기 때문에 이 부분은 넘어가고

이 문제에서 중요한 것은 DFS + DP를 이용하여 정점의 서브트리 개수를 알아내는 것이다.

 

루트는 이미 정해져있기 때문에 루트를 시작으로 각 정점과 인접한 정점을 dfs로 탐색하면서
현재 루트가 R일 때 나올 수 있는 서브트리의 개수를 배열의 메모이제이션으로 구해낸다.

 

1. 주어진 간선의 정보를 모두 저장한다.

for (int i = 0; i < N - 1; i++)
	{
		int U, V;
		cin >> U >> V;

		link[U].push_back(V);
		link[V].push_back(U);
	}

 

2. 인접한 정점을 dfs로 탐색하며 메모이제이션 과정을 거친다.

int dfs(int root)
{
	if (subtree[root] != 0) return subtree[root];
	visited[root] = true;
	int sub = 1;

	for (int i = 0; i < link[root].size(); i++)
	{
		int next = link[root][i];
		if (!visited[next])
		{
			sub += dfs(next);
		}
	}
	subtree[root] = sub;
	return sub;
}

루트를 기준으로 정점의 끝까지 탐색 후, 역으로 값을 메모이제이션 하며 정점마다의 
서브트리 개수를 저장해 나가는 과정을 거친다.

풀이

#include <iostream>
#include <algorithm>
#include <vector>
#include <cmath>
#include <string>
#include <queue>
#include <string.h>
#include <limits.h>
#include <cstdio>
#include <map>

using namespace std;

vector<vector<int>> link(100001);
bool visited[100001];
int subtree[100001];

int dfs(int root)
{
	if (subtree[root] != 0) return subtree[root];
	visited[root] = true;
	int sub = 1;

	for (int i = 0; i < link[root].size(); i++)
	{
		int next = link[root][i];
		if (!visited[next])
		{
			sub += dfs(next);
		}
	}
	subtree[root] = sub;
	return sub;
}

int main()
{
	ios::sync_with_stdio(false);
	cin.tie(0);
	cout.tie(0);
	
	int N, R, Q;

	cin >> N >> R >> Q;

	for (int i = 0; i < N - 1; i++)
	{
		int U, V;
		cin >> U >> V;

		link[U].push_back(V);
		link[V].push_back(U);
	}

	subtree[R] = dfs(R);

	for (int i = 0; i < Q; i++)
	{
		int n;
		cin >> n;

		cout << subtree[n] << '\n';
	}
}