문제
N개의 정점으로 이루어진 트리에서 루트 R이 주어진다. Q개의 쿼리마다 특정 정점을 루트로 하는 서브트리의 정점 수를 출력하라.
입력
첫째 줄에 N, R, Q가 주어진다. 이후 N-1개의 간선, Q개의 쿼리 정점이 주어진다.
출력
각 쿼리에 대해 서브트리 크기를 출력한다.
예제
| 입력 | 출력 |
|---|---|
9 5 3 1 3 4 3 5 4 5 6 6 7 2 3 9 6 8 6 5 4 8 | 9 4 1 |
풀이
DFS로 루트부터 서브트리 크기를 메모이제이션한 뒤, 각 쿼리를 O(1)에 응답한다.
- 루트 R에서 DFS를 시작하여 각 노드의 서브트리 크기를 재귀적으로 계산한다
dp[n] = 1 + sum(dp[자식])으로 자신을 포함한 서브트리 크기를 저장한다- 부모 노드를 매개변수로 전달하여 역방향 탐색을 방지한다
- 각 쿼리는 미리 계산된
dp[u]값을 출력한다
핵심 아이디어: DFS 한 번으로 모든 서브트리 크기를 전처리하면 Q개의 쿼리를 각각 O(1)에 처리할 수 있다.
코드
import java.io.*;
import java.util.*;
public class Day364BOJ15681트리와쿼리 {
static int N, R, Q;
static Integer[] dp;
static List<Integer>[] list;
public static void main(String[] args) throws Exception {
BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
StringTokenizer st = new StringTokenizer(br.readLine());
StringBuilder sb = new StringBuilder();
N = Integer.parseInt(st.nextToken());
R = Integer.parseInt(st.nextToken()) - 1;
Q = Integer.parseInt(st.nextToken());
dp = new Integer[N];
list = new ArrayList[N];
for (int i = 0; i < N; i++)
list[i] = new ArrayList<>();
for (int i = 0; i < N - 1; i++) {
st = new StringTokenizer(br.readLine());
int u = Integer.parseInt(st.nextToken()) - 1;
int v = Integer.parseInt(st.nextToken()) - 1;
list[u].add(v);
list[v].add(u);
}
recur(R, -1);
for (int i = 0; i < Q; i++)
sb.append(dp[Integer.parseInt(br.readLine()) - 1]).append("\n");
System.out.println(sb);
br.close();
}
private static int recur(int n, int p) {
if (dp[n] == null) {
dp[n] = 1;
for (int i : list[n])
if (p != i)
dp[n] += recur(i, n);
}
return dp[n];
}
}복잡도
- 시간: O(N + Q) - DFS 전처리 O(N) + 쿼리 O(Q)
- 공간: O(N)