baekjoon

[백준] 15681 트리와 쿼리

윤만석 2023. 3. 31. 10:47

문제

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

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

 

원래트리는 아무 노드나 잡고 루트를 잡을 수 있지만 이 문제에서는 루트를 정해주고, 

어떤 노드 N의 서브트리의 노드개수를 물어봤습니다.

 

따라서, 탑다운방식으로 DP를 사용하여 루트를 기준으로 리프노드부터 개수를 세서 저장해주면 됩니다.

#include<bits/stdc++.h>
#define FAST ios_base::sync_with_stdio(false),cin.tie(NULL);
#define mset(v) memset(v,0,sizeof(v));
#define rep(i,a) for(int i=0;i<a;++i)
#define REP(i,a) for(int i=1;i<=a;++i)

using namespace std;

typedef long long ll;
typedef pair<int, int> pi;
typedef tuple<int, int, int>ti;
typedef vector<int> vi;
typedef vector<vector<int>> vvi;
int dy[] = { -1,0,1,0 }, dx[] = { 0,1,0,-1 };
int N, R, Q, U, V, X, v[100001], dp[100001];
vi adj[100001];

int init(int k) {
	v[k] = 1;
	dp[k] = 1;
	for (auto c : adj[k]) {
		if (v[c])continue;
		init(c);
		dp[k] += dp[c];
	}
	return dp[k];
}
int main() {
	FAST;
	cin >> N >> R >> Q;
	rep(i, N - 1) {
		cin >> U >> V;
		adj[U].push_back(V);
		adj[V].push_back(U);
	}
	init(R);
	while (Q--) {
		cin >> X;
		cout << dp[X] << "\n";
	}
}