[BOJ] 15681 : 트리와 쿼리

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

---

문제해석

 루트가 정해진 트리가 주어졌을 때, 임의의 노드의 SubTree의 크기를 구하자.

 이때, 트리가 이진트리! 라는 말이 없으므로 조심해야한다.

문제핵심

 이 문제의 핵심은 SubTree의 크기를 구해야되는 쿼리가 10^5번 들어올 수 있기 때문에, 처음에 Root부터 시작해서 모든 노드의 SubTree의 크기를 구해둬야한다. ( 그렇게 안해도 시간 초과가 안나오는지는 모르겠습니다! ) 쿼리마다 구하면 아마 시간초과가 나올 것 같다.

알고리즘

 TreeDp를 이용해서 Root 부터 시작해서 재귀적으로 모든 노드들의 SubTree 의 크기를 구해야한다. 현재 노드가 연결된 노드들이 이미 다 방문한 노드라면, 해당 노드는 leaf 노드 이므로 Dp 값은 0이고, 1을 return 해주면 된다. leaf 노드가 아니라면 자식 노드들의 subTree 값을 모두 더해서 저장하고, 자기 자신을 포함해서 subTree + 1 값을 return 해주면 된다.

#include<iostream>
#include<vector>
using namespace std;
int subtree[100001];
vector<int> tree[100001];
bool visited[100001];
int countSubtree(int cur) {
    // 연결된 애들 중에 더 이상 체크할게 없을 때 --> leaf 노드일 때 1
    bool check = false;
    int sum = 0;
    for (int i = 0; i < tree[cur].size(); i++) {
        int next = tree[cur][i];
        if (!visited[next]) {
            visited[next] = 1;
            sum += countSubtree(next);
            check = true;
        }
    }
    if (!check) return 1;
    subtree[cur] = sum;
    return subtree[cur] + 1;
}
int main(void) {
    ios::sync_with_stdio(false);
    cin.tie(0); cout.tie(0);
    int N, R, Q, u, v;
    cin >> N >> R >> Q;
    for (int i = 0; i < N - 1; i++) {
        cin >> u >> v;
        tree[u].push_back(v);
        tree[v].push_back(u);
    }
    visited[R] = 1;
    countSubtree(R);
    for (int i = 0, x; i < Q; i++) {
        cin >> x;
        cout << subtree[x] + 1 << '\n';
    }
    return 0;
}

---