[C/C++] 백준 15681번 - 트리와 쿼리

✏️ 문제 링크

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

15681번: 트리와 쿼리

 

✏️ 문제 설명

 

 

✏️ 문제 풀이

트리를 구현하고 서브트리의 개수를 구하는 문제입니다.

각 노드의 부모를 저장할 parent배열을 만든 후 parent[루트]의 값을 -1로 설정해줍니다.

 

루트번호에서 dfs함수를 돌려서 트리를 구성해주는데 저는 여기서 조건을 두 가지를 걸었습니다.

간선이 무방향으로 주어지다 보니 부모 노드로 갈 수 있는 경우도 생길 수 있기 때문에 parent[node]== -1(루트)이거나

parent[node]의 값이 0이 아니라면(이미 부모 노드가 있다는 것임) 방문하지 않도록 설계했습니다.

 

서브트리의 개수를 구하는 과정입니다. 처음에는 입력받는 정점에 대해서 모두 dfs를 돌려서 서브트리의 개수를 구할려고 했는데 52%에서 시간 초과가 발생하였습니다. 문제의 힌트를 다시 한 번 읽어보니 루트번호에서 dfs를 한 번만 돌려서 각 정점의 서브트리의 개수를 구한 뒤 입력받는 정점에 대해서는 size[u]만 출력하게 할 수 있다는 사실을 알았습니다.

N+1크기의 size배열을 선언 후, dfs를 돌리면서 자신이 방문하는 노드의 사이즈를 1로 설정해주고, 나중에 부모 노드에 대해서는 자식 노드의 크기를 더해주는 방식으로 구현해줍니다.

아마 아래 코드를 보시면 좀 더 이해가 잘 될것 같습니다.

 

 

✏️ 문제 풀이

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

vector<vector<int>>graph;
vector<int>parent;
vector<int>sevSize;
int N, R, Q;
void makeTree(int route);
void countNode(int node);
int count1 = 0;

int main() {

	ios::sync_with_stdio(false);
	cin.tie(NULL);
	cout.tie(NULL);
	cin >> N >> R >> Q;

	graph.resize(N + 1);
	parent.resize(N + 1);
	sevSize.resize(N + 1);

	parent[R] = -1;

	for (int i = 0; i < N-1; i++) {
		int u, v;
		cin >> u >> v;
		graph[u].push_back(v);
		graph[v].push_back(u);
	}
	makeTree(R);
	countNode(R);
	for (int i = 0; i < Q; i++) {
		int node;
		cin >> node;
		if (node == R)
			cout << N << "\n";
		else {			
			cout << sevSize[node] << "\n";
		}
	}

}
void makeTree(int route) {
	for (int next : graph[route]) {
		if (parent[next] == -1 || parent[next] != 0)
			continue;
		parent[next] = route;
		makeTree(next);
	}
}
void countNode(int node) {
	sevSize[node] = 1;
	for (int next : graph[node]) {
		if (parent[node] == next)
			continue;
		countNode(next);
		sevSize[node] += sevSize[next];
	}
}