[백준/C++] 15681번: 트리와 쿼리

문제

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

---

문제 요약

간선에 가중치와 방향성이 없는 임의의 루트(R)가 있는 트리가 주어졌을 때, 각 쿼리에 대한 답을 출력한다.

 

쿼리는 정점 U가 주어지고, U를 루트로 하는 서브트리에 속한 정점의 수를 출력한다.

---

코드

#include <iostream>
#include <algorithm>
#include <vector>

using namespace std;

int n, r, q;

vector<int> edges[100010];

// subtree[i]: i번 노드를 루트로 하는 서브트리에 속한 노드 수
int subtree[100010];
bool visited[100010];

// x번 노드를 루트로 하는 서브트리에 속한 노드 수를 반환하는 함수
int dfs(int x) {
    visited[x] = true;

    for (int i = 0; i < edges[x].size(); i++) {
        int y = edges[x][i];

        if (visited[y]) continue;

        // 자식 노드를 루트로 하는 서브트리에 속한 노드 수 더하기
        subtree[x] += dfs(y);
    }

    // +1은 x번 노드 자기 자신을 포함
    return subtree[x] += 1;

}

int main() {
    // 입출력 시간 단축
    ios::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);

    cin >> n >> r >> q;

    // 그래프 정보 입력
    for (int i = 0; i < n - 1; i++) {
        int a, b;
        cin >> a >> b;

        edges[a].push_back(b);
        edges[b].push_back(a);
    }

    // DFS 탐색
    dfs(r);

    // 쿼리 처리
    for (int i = 0; i < q; i++) {
        int num;
        cin >> num;

        cout << subtree[num] << "\n";
    }

    return 0;
}

---

코드 설명

1. 그래프 정보를 입력 받는다. (양방향)

 

2. R번 루트를 시작으로 DFS 탐색을 진행하며, 자식 노드를 루트로 하는 서브트리에 속한 노드 수를 구한다.

+ dfs 함수는 인자로 받은 x번 노드를 루트로 하는 서브 트리에 속한 노드 수를 반환한다.

 

3. 이 때, 1차원 배열 subtree[x]에 반환 값을 저장한다.

 

4. 각 쿼리에 대해 subtree 값을 출력한다.

 

x번 노드를 루트로 하는 서브 트리에 속한 노드 수

= x번 노드의 자식들을 루트로 하는 각 서브 트리에 속한 노드 수