[BOJ] 15681 - 트리와 쿼리 (C++)

문제

트리와 쿼리

코드

#include <iostream>
#include <vector>

using namespace std;

int N, R, Q;
vector<int> e[100001];
bool visited[100001];
int nums[100001];

int dfs(int r){
    if (nums[r] != 0) return nums[r];

    visited[r] = true;

    int cnt = 1;
    for (int i = 0; i < e[r].size(); i++){
        int next = e[r][i];
        if (visited[next]) continue;
        cnt += dfs(next);
    }
    nums[r] = cnt;
    return cnt;
}

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

    cin >> N >> R >> Q;
    
    int u, v;
    for (int i = 0; i < N-1; i++){
        cin >> u >> v;
        e[u].push_back(v);
        e[v].push_back(u);
    }

    nums[R] = dfs(R);

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

접근

각 노드를 기준으로 그 노드를 루트로 하는 서브 트리에 속한 노드의 개수가 몇 개인지 저장하려고 했다. 이를 위해 주어진 루트노드를 기준으로 dfs를 수행하였고 이때 부모 노드는 개수에 들어가지 않게 하기 위해 visited 배열을 사용하였다. 방향성이 없기 때문에 간선 정보를 노드 양쪽에 저장한 것과 visited 배열로 중복 탐색을 고려한 부분이 핵심이었던 것 같다.