문제 출처

1. 문제설명

  • 요약하자면 다음과 같다.
    • 어떤 그래프와 시작점이 주어질 때, 각 노드의 서브 트리 개수를 구하라
  • 문제 분류
    • Graph, Tree, DFS, DP in Tree


2. 알고리즘 설계

  • 트리의 성질
    • 무방향 그래프
    • 사이클 없음.
    • child node가 없어도 트리이다.
  • 로직
    • 인접 행렬에 간선을 저장한다. (무방향 이므로, 양쪽 방향 모두 저장)
    • 시작점 R을 인자로 DFS 실행
    • 트리는 child가 없어도 트리이기 때문에 현재 정점의 서브트리 개수를 1로 저장
    • 연결된 정점을 하나씩 꺼내면서, 연결된 서브트리의 개수를 더함.
    sub_tree[cur] = 1;  // 반드시 첫줄에 선언 혹은 배열 원소 전부 1로 미리 초기화

for(int nxt : adj[cur]) {
  if(sub_tree[nxt] != 0) continue;
  sub_tree[cur] += count(nxt);  // child의 서브트리 개수를 더한다.
}

return sub_tree[cur];


3. 전체 코드

#include <bits/stdc++.h>
using namespace std;
using pii = pair<int, int>;

const int SZ = 100'002;
int N, R, Q, U, V;
vector<int> adj[SZ];
int vis[SZ];
int sub_tree[SZ];

int count(int cur) {
    sub_tree[cur] = 1;
    
    for(int nxt : adj[cur]) {
        if(sub_tree[nxt] != 0) continue;
        sub_tree[cur] += count(nxt);
    }

    return sub_tree[cur];
}

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

    cin >> N >> R >> Q;
    for(int i=1; i<N; i++) {
        cin >> U >> V;
        adj[U].push_back(V);
        adj[V].push_back(U);
    }

    count(R);

    while(Q--) {
        cin >> U;
        cout << sub_tree[U] << '\n';
    }

    return 0;
}


4. 참고

  • 이 문제부터 문제 분류를 추가하였다.