백준 15681 - 트리와 쿼리

주어진 입력으로 트리를 구성한다.

 

그리고 그래프 탐색을 통해 해당 노드의 깊이도 저장하도록 한다.

 

그 후 탑다운 DP를 통해 해당 노드가 가지는 자식의 갯수를 계산하여 출력하도록 하면 된다.

 

전체 코드

#include <bits/stdc++.h>

using namespace std;

vector<int> al[100001];
int depth[100001];
int dp[100001];

int build_tree(int v);

int main()
{
    memset(dp, -1, sizeof dp);
    ios::sync_with_stdio(false); cin.tie(0);
    
    int v, rt, qry;
    cin >> v >> rt >> qry;
    
    for (int i = 0; i < v - 1; ++i)
    {
        int go, to;
        cin >> go >> to;
        al[go].push_back(to);
        al[to].push_back(go);
    }

    queue<pair<int, int>> q;
    q.push({rt, 1});
    
    while (q.size())
    {
        auto cur = q.front();
        q.pop();
        
        depth[cur.first] = cur.second;
        
        for (auto item : al[cur.first])
        {
            if (depth[item])
                continue;
            
            q.push({item, cur.second + 1});
        }
    }
    
    build_tree(rt);
    
    while (qry--)
    {
        int node;
        cin >> node;
        cout << dp[node] + 1 << '\n';
    }
}

int build_tree(int v)
{
    int &ret = dp[v];
    
    if (ret != -1)
        return ret;
    
    ret = 0;
    
    for (int item : al[v])
        if (depth[v] + 1 == depth[item])
            ret += 1 + build_tree(item);
        
    return ret;
}