[백준,c++] 15681번 - 트리와 쿼리

문제

15681번: 트리와 쿼리

 

코드

#include <iostream>
#include <vector>

using namespace std;

int N, R, Q;
int max_depth = -1, max_node;
vector<int>graph[100001];
int subnode_cnt[100001];

int dfs(int node, int before){
    for(auto next: graph[node]){
        if(next != before){
            subnode_cnt[node] += dfs(next, node)+1;
        }
    }
    return subnode_cnt[node];
}

int main(){
    ios_base::sync_with_stdio(false);
    cin.tie(0);
    cin>>N>>R>>Q;
    for(int i=0; i<N-1; i++){
        int a,b; cin>>a>>b;
        graph[a].push_back(b);
        graph[b].push_back(a);
    }
    dfs(R, 0);
    
    for(int i=0; i<Q; i++){
        int q; cin>>q;
        cout<<subnode_cnt[q]+1<<'\n';
    }
}

 

풀이

해당 노드를 포함하여 자식 노드의 개수가 몇 개인지를 출력하는 문제입니다.

노드의 수가 굉장히 많기 때문에 dfs를 한번만 돌아서 노드의 개수를 파악해야 합니다.

핵심 로직은  subnode_cnt[node] += dfs(next, node)+1;  입니다.

리프 노드까지 들어간 다음 루트 노드로 돌아오면서 노드의 개수를 +1씩 메모이제이션 함으로써 단 한번은 dfs만으로 subnode의 갯수를 파악할 수 있습니다.