[BOJ / C++] 트리와 쿼리 (트리)

문제링크: https://www.acmicpc.net/problem/15681

 

1. 코드

#include <bits/stdc++.h>
using namespace std;

void dfs(int cur, int parent, vector<vector<int>>& adj, vector<int>& subtree) {
    subtree[cur] = 1;
    for (int nxt : adj[cur]) {
        if (nxt == parent) continue;
        dfs(nxt, cur, adj, subtree);
        subtree[cur] += subtree[nxt];
    }
}

int main() {
    ios::sync_with_stdio(0); cin.tie(0);

    int n, r, q;
    cin >> n >> r >> q;

    vector<vector<int>> adj(n+1);
    vector<int> parent(n+1, -1);

    for (int i = 0; i < n-1; i++) {
        int u, v;
        cin >> u >> v;

        adj[u].push_back(v);
        adj[v].push_back(u);
    }

    // 서브트리 테이블 채우기
    vector<int> subtree(n+1, 0);
    dfs(r, -1, adj, subtree);

    // query 진행
    while (q--) {
        int root;
        cin >> root;

        cout << subtree[root] << '\n';
    }
}

 

 

2. 시간복잡도

\(n=10^{5}, q=10^{5}\)이다. 만약 쿼리당 BFS를 해서 search할 경우 \(O(nq)=10^{10} > 10^{8}\)으로 시간초과가 된다. 그럼 미리 subtree table을 채워두면 되는가 하면 그것도 아니다. 모든 노드에 대해 search를 수행해야하므로 \(O(n^2)=10^{10} > 10^{8}\)으로 여전히 시간초과가 된다. 즉, \(n=10^{5}\)으로 꽤 크므로 한번의 search에 subtree table을 채울 수 있어야한다.

 

핵심은 DFS를 이용해서 bottom up 방식으로 divide and conquer를 진행한다. DFS를 이용해 leaf node까지 타고 내려가서 subtree[leafnode] = 1로 하고 위로 올라오면서 자식노드의 subtree 개수를 더해주면 된다. 이런 방식을 이용하면 한번의 DFS 탐색으로 \(O(n)=10^{5} < 10^{8}\)에 subtree 테이블을 전부 채울 수 있다.

 

3. 분석

 

문제 아래에 자세히 설명되어있긴 한데 먼저 서브트리가 뭔지를 알아야한다. 서브트리는 부모 노드를 끊어냈을 때 해당 정점과 그 자식으로 구성된 트리를 말한다. 해당 정점에서의 서브트리의 노드 개수를 알아내려면 해당 정점에서 자식노드쪽으로 DFS 또는 BFS로 탐색을 해서 count를 해야한다. 부모로는 DFS, BFS가 들어가면 안되기 때문에 부모에 대한 정보는 들고 있어야한다.

 

트리에서 간선의 개수는 정점개수 -1 인것도 주어지긴했지만 기본적으로 알고 있어야한다.