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

15681번 트리와 쿼리

문제 요약

문제

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

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

사용 알고리즘

그래프, DFS

입출력

• 입력: 트리의 정점의 수 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줄에 걸쳐 각 쿼리의 답을 정수 하나로 출력한다. ``` 예제1

입력9 5 3 1 3 4 3 5 4 5 6 6 7 2 3 9 6 6 8 5 4 8 출력 9 4 1

## 풀이
DFS로 쿼리를 확인할때마다 하는 것이 아닌 미리 모든 노드의 서브트리의 정점 개수를 구해놓는다. 먼저 DFS로 도착한 노드에 정점개수를 1로 잡아 놓는다. 자기자신도 포함하기 때문이다. 이후 DFS를 통해 온 부모를 제외한 연결된 모든 자식 노드에 재귀함수로 반복하여 개수를 구하고 자신에게 추가한다. 만약 자식이 없거나 모든 자식을 보고 왔다면 부모에게 자신의 개수를 보낸다.

## 어려웠던 점
DFS를 할때 재귀함수에 어디서 왔는지에 대한 정보를 넣어서 부모를 제외시키고 돌려야 한다. 나머지는 DFS로 미리 처리해 놓으면 쉽게 풀 수 있다.

## 전체 코드
```cpp
#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

int n, r, q, u, v;
vector<vector<int>> nodes;
vector<int> count_node;

int count(int pre, int now){
    count_node[now] = 1;
    for (int i = 0; i < nodes[now].size(); i++) {
        if(nodes[now][i] != pre){
            count_node[now] += count(now, nodes[now][i]);
        }
    }
    return count_node[now];
}


int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cin >> n >> r >> q;
    nodes = vector<vector<int>>(n+1);
    count_node = vector<int>(n+1, 0);
    for (int i = 0; i < n-1; i++) {
        cin >> u >> v;
        nodes[u].push_back(v);
        nodes[v].push_back(u);
    }
    count(-1, r);
    for (int i = 0; i < q; i++) {
        cin >> u;
        cout << count_node[u]<< endl;
    }
    
    
}