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

코드

#include <iostream>
#include <algorithm>
#include <vector>

using namespace std;

int n,r,q,u,v,k,res, visited[100003],dp[100003];
vector<int> adj[100003];

// dfs로 그래프를 순회하면서, dp 배열에 값 저장 => 서브트리에 속한 정점의 수 저장

int dfs(int x) {
    if(dp[x] != 0) return dp[x]; // 메모이제이션
    visited[x] = 1;
    int ret = 1;
    for(int i : adj[x]) {
        if(!visited[i]) {
            ret += dfs(i);
        }
    }
    dp[x] = ret;
    return ret;
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);
    cin >> n >> r >> q;
    // 양방향 그래프 생성
    for(int i = 0; i < n - 1; i++) {
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    dp[r] = dfs(r);
    for(int i = 0; i < q; i++) {
        cin >> k;
        cout << dp[k] << '\n';
    }
    
    return 0;
}

풀이

dfs로 그래프를 순회하면서, dp 배열에 아래 노드가 몇 개인지 저장한다. 전형적인 top-bottom 문제인 듯하다. 처음에 감을 못잡아서 헤맸다. 그 후 루트 노드를 지정하고, 트리를 순회한다. 그리고 입력값에 따라 적절한 답을 출력하면 된다.