[BOJ 15681] 트리와 쿼리

문제

www.acmicpc.net/problem/15681

15681번: 트리와 쿼리

 

---

풀이

트리DP를 사용합니다.

트리가 주어졌을 때, 배열 \(dp\)를 다음과 같이 정의할 수 있습니다:

 

\(dp[i]:\) \(i\)를 루트로 하는 서브트리에 속한 정점의 수.

 

트리를 순회하면서 \(i\)의 자식 노드들을 \(c_{1}, c_{2}, c_{3}, ... , c_{m}\)이라고 할 때, \(dp[i]\)에 다음 값을 저장합니다:

$$dp[i] = 1 + \sum_{j = 1}^{m} dp[c_{j}]$$

각 쿼리마다 \(U\)가 주어졌을 때 \(dp[U]\)를 출력하면 됩니다.

 

---

코드

#include <bits/stdc++.h>
using namespace std;

int N, R, Q;
vector<vector<int>> arr;
vector<int> dp;

int solve(int cur, int pnode) {
    for (auto node : arr[cur])
        if (node != pnode)
            dp[cur] += solve(node, cur);

    return dp[cur];
}

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

    cin >> N >> R >> Q;
    arr.resize(N + 1);
    dp.resize(N + 1, 1);

    for (int i = 0; i < N - 1; ++i) {
        int a, b;
        cin >> a >> b;
        arr[a].push_back(b);
        arr[b].push_back(a);
    }

    solve(R, 0);

    while (Q--) {
        int query;
        cin >> query;
        cout << dp[query] << "\n";
    }

    return 0;
}