baekjoon 15681:트리와 쿼리

baekjoon 15681 트리와 쿼리

15681번 트리와 쿼리

접근

dfs를 통해 트리를 만든 후 dp for tree를 이용해서 count를 센다.

코드

#include <iostream>
#include <vector>
using namespace std;
int n, r, q;
void dfsTree(int x, vector<vector<int>>& adj, vector<vector<int>>& tree, vector<bool>& visited){
  visited[x] = true;
  for(int i = 0; i < adj[x].size(); i++){
    if(!visited[adj[x][i]]){
      tree[x].push_back(adj[x][i]);
      dfsTree(adj[x][i], adj, tree, visited);
    }
  }
};
int dfsCnt(int x, vector<vector<int>>& tree, vector<int>& count){
  int ret = 1;
  for(int i = 0; i < tree[x].size(); i++){
    if(!count[tree[x][i]]){
      ret += dfsCnt(tree[x][i], tree, count);
    }
    else{
      ret += count[tree[x][i]];
    }
  }
  return count[x] = ret;
};

int main(){
  ios::sync_with_stdio(false);
  cin.tie(NULL); cout.tie(NULL);
  cin >> n >> r >> q;
  vector<vector<int>> adj(n);
  vector<vector<int>> tree(n);
  vector<int> count(n);
  vector<bool> visited(n, false);
  for(int i = 0; i < n-1; i++){
    int a, b; cin >> a >> b;
    adj[a-1].push_back(b-1);
    adj[b-1].push_back(a-1);
  }
  dfsTree(r-1, adj, tree, visited);
  dfsCnt(r-1, tree, count);
  
  for(int i = 0; i < q; i++){
    //query
    int x; cin >> x;
    cout << count[x-1] << "\n";
  }

  return 0;
}

배운 점

dp for tree은근 자주 쓰이는 것 같다.