[백준] Baekjoon Online Judge
문제
간선에 가중치와 방향성이 없는 임의의 루트 있는 트리가 주어졌을 때, 아래의 쿼리에 답해보도록 하자.
- 정점 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줄에 걸쳐 각 쿼리의 답을 정수 하나로 출력한다.
풀이
#include <iostream>
#include <vector>
using namespace std;
#define MAX 100001
int N, R, Q;
int cnt[MAX];
vector<int> edge[MAX];
void dfs(int now, int parent);
void dfs(int now, int parent) {
cnt[now] = 1;
for (int node : edge[now]) {
if (cnt[node] == 0) {
dfs(node, now);
cnt[now] += cnt[node];
}
}
}
int main() {
ios::sync_with_stdio(false);
cin.tie(0);
cin >> N >> R >> Q;
for (int i = 1; i < N; i++) {
int a, b;
cin >> a >> b;
edge[a].push_back(b);
edge[b].push_back(a);
}
dfs(R, 0);
while (Q--) {
int U;
cin >> U;
cout << cnt[U] << '\n';
}
return 0;
}
for (int i = 1; i < N; i++) {
int a, b;
cin >> a >> b;
edge[a].push_back(b);
edge[b].push_back(a);
}
간선을 입력 받는 부분의 코드이다. 두 정점 중 어느 것이 부모 노드인지 알 수 없기에 해당 간선에 대한 정보를 두 노드에 모두 저장하도록 했다.
void dfs(int now, int parent) {
cnt[now] = 1; // 자기자신을 포함해야 하기에
for (int node : edge[now]) {
if (cnt[node] == 0) {
dfs(node, now);
cnt[now] += cnt[node];
}
}
}
깊이우선탐색으로 각 노드의 자식의 수를 구하는 함수이다. 부분트리의 정점의 수에는 자기자신을 포함해야 하기에 함수의 시작부분에 cnt[now]를 1로 초기화하도록 했다. 그 후 now 노드와 간선으로 연결된 노드들을 탐색하도록 했는데, 이때 cnt[]가 0인 방문한 적이 없는 노드만 방문하게 된다.
'백준 > C++' 카테고리의 다른 글
[Baekjoon/C++] 1197번 - 최소 스패닝 트리 (크루스칼 알고리즘) (0) | 2024.07.30 |
---|---|
[Baekjoon/C++] 31611번 - 火曜日 (Tuesday) (0) | 2024.07.29 |
[Baekjoon/C++] 17404번 - RGB거리 2 (0) | 2024.07.25 |
[Baekjoon/C++] 1644번 - 소수의 연속합 (0) | 2024.07.21 |
[Baekjoon/C++] 9252번 - LCS 2 (2) | 2024.07.20 |