[BOJ] 15861 - 트리와 쿼리

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

BOJ - 트리와 쿼리

문제 유형: 트리, dfs

문제 난이도: Gold V

시간 제한: 1초

메모리 제한: 128MB

 

문제

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

• 정점 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줄에 걸쳐 각 쿼리의 답을 정수 하나로 출력한다.

 

풀이

dfs 한 번으로 모든 서브트리의 수를 계산해둘 수 있다.

서브트리크기 = 자식서브트리크기들의 합 + 1(루트)이다.

이를 이용해서 dfs를 구현해주면 된다.

 

각 쿼리에 대한 응답을 해주면 된다.

 

코드

C++

#include <bits/stdc++.h>

using namespace std;

int n, r, q;
void solve(int root, vector<bool> &visited, vector<vector<int>> &graph,
           vector<int> &subtreesize) {
  for (const auto &child : graph[root]) {
    if (!visited[child]) {
      visited[child] = true;
      solve(child, visited, graph, subtreesize);
      subtreesize[root] += subtreesize[child];
    }
  }
  subtreesize[root]++;
}
int main(int argc, char *argv[]) {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  cout.tie(0);
  cin >> n >> r >> q;
  vector<vector<int>> graph(n + 1);

  for (int i = 0; i < n - 1; ++i) {
    int u, v;
    cin >> u >> v;
    graph[u].push_back(v);
    graph[v].push_back(u);
  }

  vector<int> subtreesize(n + 1, 0);
  vector<bool> visited(n + 1, false);
  visited[r] = true;
  solve(r, visited, graph, subtreesize);

  for (int i = 0; i < q; ++i) {
    int root;
    cin >> root;
    cout << subtreesize[root] << "\n";
  }

  return 0;
}