트리와 쿼리
문제 해석
루트 번호와 간선 정보가 주어졌을 때, U를 루트로 하는 서브트리에 속한 정점의 수를 출력하는 문제입니다.
설계(10)
지문에 힌트가 아주 친절하게 나와있습니다.
우선, 주어진 루트를 기준으로 dfs탐색을 이용하여 tree를 만듭니다.
이와 같이 부모 노드, 자식노드를 갱신해줍니다.
그 후, 루트를 시작으로 하여 서브트리의 정점 개수를 카운트 합니다.
시간복잡도는 트리 생성 O(N) + 서브트리의 정점의 개수 카운트 O(N)으로 O(N)으로 해결이 가능합니다.
구현(10)
코드
#include <iostream>
#include <vector>
using namespace std;
#define MAXN 100001
int len[MAXN];
struct Tree{
int par=-1;
vector<int> children;
};
int N,R,Q;
vector<int> adj[MAXN];
Tree tree[MAXN];
void makeTree(int root, int par){
for (int i = 0;i<adj[root].size(); i++){
int next = adj[root][i];
if (next != par){
tree[root].children.push_back(next);
tree[next].par = root;
makeTree(next,root);
}
}
}
void countSubtreeNodes(int root){
len[root] = 1;
for (auto x : tree[root].children){
countSubtreeNodes(x);
len[root] += len[x];
}
}
int main(){
ios::sync_with_stdio(false);
cin.tie(0);
cin >> N >> R >> Q;
for (int i = 0 ;i <N-1; i++){
int u,v;
cin >> u >> v;
adj[u].push_back(v);
adj[v].push_back(u);
}
makeTree(R,-1);
countSubtreeNodes(R);
for (int i = 0 ;i <Q; i++){
int num;
cin >> num;
cout << len[num]<<"\n";
}
return 0;
}
디버깅
제출결과
84 ms
마무리
기본적인 트리 문제였습니다. 메모이제이션을 이용하면 트리를 만들지 않아도 충분히 구현 가능할 것으로 보입니다.