Baekjoon 15681번 트리와 쿼리

문제

https://www.acmicpc.net/problem/15681

걸린 시간

00 : 39 : 33

풀이

C++

#include <bits/stdc++.h>
#define INF 1e9
#define all(c) c.begin(), c.end()
typedef long long ll;
typedef unsigned long long ull;
using namespace std;

int n, r, q;
int u, v;
vector<set<int>> adj;
vector<int> query;

void makeTree(){
    queue<int> q;
    q.push(r);
    while(!q.empty()){
        int v = q.front();
        q.pop();
        for(auto it = adj[v].begin(); it != adj[v].end(); it++){
            int u = *it;
            adj[u].erase(v);
            q.push(u);
        }
    }
}

int count(int v){
    // 기저 사례
    if(adj[v].empty()) return 0;
    // 재귀 호출
    int& ret = query[v];
    for(auto it = adj[v].begin(); it != adj[v].end(); it++){
        int u = *it;
        ret += count(u)+1;
    }
    return ret;
}

int main(){
    ios::sync_with_stdio(false);
    cin.tie(NULL);
    cin >> n >> r >> q;
    adj.resize(n+1);
    query.resize(n+1, 0);
    for(int i = 0; i < n-1; i++){
        cin >> u >> v;
        adj[u].insert(v);
        adj[v].insert(u);
    }
    makeTree();
    count(r);
    for(int i = 0; i < q; i++){
        cin >> u;
        cout << query[u]+1 << "\n";
    }
    return 0;
}

주어진 정점 r 을 루트로 하는 트리 형태를 만들고, 각 정점을 기준으로 서브트리에 속한 정점의 개수를 분할정복으로 구한 뒤 쿼리에 맞게 값을 출력하면 되는 문제이다.

구현 과정은 어렵지 않았는데 각 쿼리마다 주어지는 정점 u 를 루트로 하는 트리에서 서브트리에 속한 정점의 개수를 구하는 문제로 잘못 이해하는 바람에 시간이 조금 걸렸다.