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

문제 15681번: 트리와 쿼리

간선에 가중치와 방향성이 없는 임의의 루트 있는 트리가 주어졌을 때, 아래의 쿼리에 답해보도록 하자.

• 정점 U를 루트로 하는 서브트리에 속한 정점의 수를 출력한다.

만약 이 문제를 해결하는 데에 어려움이 있다면, 하단의 힌트에 첨부한 문서를 참고하자.

입력

트리의 정점의 수 N과 루트의 번호 R, 쿼리의 수 Q가 주어진다. (2 ≤ N ≤ 105, 1 ≤ R ≤ N, 1 ≤ Q ≤ 105)

이어 N-1줄에 걸쳐, U V의 형태로 트리에 속한 간선의 정보가 주어진다. (1 ≤ U, V ≤ N, U ≠ V)

이는 U와 V를 양 끝점으로 하는 간선이 트리에 속함을 의미한다.

이어 Q줄에 걸쳐, 문제에 설명한 U가 하나씩 주어진다. (1 ≤ U ≤ N)

입력으로 주어지는 트리는 항상 올바른 트리임이 보장된다.

출력

Q줄에 걸쳐 각 쿼리의 답을 정수 하나로 출력한다.

문제분석

문제 요구사항

• 주어진 트리의 서브 트리 노드의 갯수를 찾기 위하여 Depth First Search가 요구됨.

이미 한번 탐색한 서브 트리 노드를 재탐색하는 비용을 절감하기 위하여 Dynamic Programming 활용.**

구현

int dp[] array를 선언하여 각 index에 해당 node를 root 로하는 서브트리의 총 노드 갯수를 저장
Initdp 함수를 통하여 dp를 완성시킴.

증명

맨 처음 dp의 모든 element들을 0으로 초기화 한 후, Initdp가 해당 dp[node]를 탐색하게 되면 dp[node]=1을 실행함, 따라서 if(dp[route]!=0) 조건문에 따라 Initdp가 탐색하지 않은 node들만 탐색함으로, Depth First Search 구현 가능.

코드

#include <iostream>
#include <vector>

std::vector<int> routes[100'001];
int dp[100'001];

const int Initdp(const int node) {
  dp[node] = 1;
  for (auto &route : routes[node]) {
    if (dp[route] != 0) {
      continue;
    }
    dp[node] += Initdp(route);
  }
  return dp[node];
}

int main() {
  std::ios::sync_with_stdio(false);
  std::cin.tie(nullptr), std::cout.tie(nullptr);

  int N, R, Q;
  std::cin >> N >> R >> Q;

  std::fill_n(dp, N + 1, 0);

  int U, V;
  for (int i = 0; i < N - 1; ++i) {
    std::cin >> U >> V;
    routes[U].push_back(V);
    routes[V].push_back(U);
  }

  Initdp(R);

  int query;
  for (int i = 0; i < Q; ++i) {
    std::cin >> query;
    std::cout << dp[query] << "\n";
  }

  return 0;
}