문제

간선에 가중치와 방향성이 없는 임의의 루트 있는 트리가 주어졌을 때, 아래의 쿼리에 답해보도록 하자.

정점 U를 루트로 하는 서브트리에 속한 정점의 수를 출력한다.
만약 이 문제를 해결하는 데에 어려움이 있다면, 하단의 힌트에 첨부한 문서를 참고하자.

입력
트리의 정점의 수 N과 루트의 번호 R, 쿼리의 수 Q가 주어진다. (2 ≤ N ≤ 105, 1 ≤ R ≤ N, 1 ≤ Q ≤ 105)
이어 N-1줄에 걸쳐, U V의 형태로 트리에 속한 간선의 정보가 주어진다. (1 ≤ U, V ≤ N, U ≠ V)
이는 U와 V를 양 끝점으로 하는 간선이 트리에 속함을 의미한다.
이어 Q줄에 걸쳐, 문제에 설명한 U가 하나씩 주어진다. (1 ≤ U ≤ N)
입력으로 주어지는 트리는 항상 올바른 트리임이 보장된다.

출력
Q줄에 걸쳐 각 쿼리의 답을 정수 하나로 출력한다.

문제 링크

풀이

해당 문제의 페이지에 힌트를 읽어보니 많은 도움이 되었다.
트리는 몇가지 특징이 있다.

• 사이클이 없다.
• 따라서 어떤 노드든 루트 노드가 될 수 있다.
• n-1 개의 간선을 가진다.
• 임의의 두 정점 U와 V에 대해, U에서 V로 가는 최단경로는 유일하다.
• 아무 정점이나 잡고 부모와의 연결을 끊었을 때, 해당 정점과 그 자식들, 그 자식들의 자식들… 로 이루어진 부분그래프는 트리가 된다. 이를 서브트리라 한다.

이 문제 처럼 노드간의 연결만 주어진 입력으로 트리를 만들기 위해서는 트리의 부모-자식 관계를 이용하면 된다.

트리에서는 자식들의 부모노드는 항상 1개이다.
어떤 정점에 대해 연결된 모든 정점은 최대 한 개의 정점을 제외하면 모두 해당 정점의 자식들이 될 것이다.
이에 따라, 부모 정점의 정보를 가져가면서,
부모 정점이 아니면서 자신과 연결되어 있는 모든 정점을 자신의 자식으로, 자신의 자식이 될 정점들의 부모 정점을 자신으로 연결한 뒤 재귀적으로 자식 정점들에게 트리 구성을 요청하는 형태의 함수이다.

void makeTree(int currentNode, int parrent) {
   	for (auto Node : graph[currentNode]) {
   		if (Node != parrent) {
   			tree[currentNode].push_back(Node);
   			makeTree(Node, currentNode);
   		}
   	}
}

makeTree를 이용해 그래프를 트리로 변경하여 저장한다.

• 부모가 -1인 경우는 부모가 없음(루트노드)을 의미한다.

트리의 정점의 개수는 1(루트노드) + 서브트리의 정점개수이다.
이때 서브트리의 정점 개수는 또 다시 1(서브트리의 루트) + 서브트리의 서브트리 정점 개수이다.
이렇게 재귀적으로 리프노드까지 간 후 그 값을 이용하여 루트 노드에서 트리의 정점의 개수를 구할 수 있게 된다.
하지만 각 노드를 루트로 하는 서브트리의 정점의 개수를 저장해두지 않으면 쿼리가 있을 때마다 불필요한 계산을 중복으로 해야한다. 따라서 배열에 각 정점이 루트노드인 서브트리의 정점의 개수를 저장한다.

void countSubtreeNodes(int currentNode) {
   	for (auto Node : tree[currentNode]) {
   		countSubtreeNodes(Node);
   		subNodes[currentNode] += subNodes[Node];
   	}
}

이 문제는 간단하여 위의 두 함수를 합쳐 하나의 함수를 만들어 풀 수 있었다.
문제가 복잡하지 않아 간단하게 구현이 가능하다.
또 새로 만들어지는 트리를 저장할 필요가 없어져 더 적은 메모리를 사용하게 된다.

코드

#include <iostream>
#include <vector>

using namespace std;

int n, r, q;
vector < vector<int>> graph;
vector<vector<int>> tree;
vector<int> subNodes;

void makeTree(int currentNode, int parrent) {
   	for (auto Node : graph[currentNode]) {
   		if (Node != parrent) {
   			tree[currentNode].push_back(Node);
   			makeTree(Node, currentNode);
   		}
   	}
}

void countSubtreeNodes(int currentNode) {
   	for (auto Node : tree[currentNode]) {
   		countSubtreeNodes(Node);
   		subNodes[currentNode] += subNodes[Node];
   	}
}

void dfs(int currentNode, int parrent){
    for (auto Node : graph[currentNode]) {
		if (Node != parrent) {
			dfs(Node, currentNode);
		    subNodes[currentNode] += subNodes[Node];
		}
	}
}

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

	cin >> n >> r >> q;
	graph.assign(n + 1, vector<int>());
	tree.assign(n + 1, vector<int>());
	subNodes.assign(n + 1, 1);
	for (int i = 0; i < n - 1; i++) {
		int a, b;
		cin >> a >> b;
		graph[a].push_back(b);
		graph[b].push_back(a);
	}

	makeTree(r, -1);
	countSubtreeNodes(r);
    dfs(r, -1);

	for (int i = 0; i < q; i++) {
		int node;
		cin >> node;
		cout << subNodes[node] << "\n";
	}

	return 0;
}