문제
간선에 가중치와 방향성이 없는 임의의 루트 있는 트리가 주어졌을 때, 아래의 쿼리에 답해보도록 하자.
- 정점 U를 루트로 하는 서브트리에 속한 정점의 수를 출력한다.
만약 이 문제를 해결하는 데에 어려움이 있다면, 하단의 힌트에 첨부한 문서를 참고하자.
입력
트리의 정점의 수 N과 루트의 번호 R, 쿼리의 수 Q가 주어진다. (2 ≤ N ≤ 10^5, 1 ≤ R ≤ N, 1 ≤ Q ≤ 10^5)
이어 N-1줄에 걸쳐, U V의 형태로 트리에 속한 간선의 정보가 주어진다. (1 ≤ U, V ≤ N, U ≠ V)
이는 U와 V를 양 끝점으로 하는 간선이 트리에 속함을 의미한다.
이어 Q줄에 걸쳐, 문제에 설명한 U가 하나씩 주어진다. (1 ≤ U ≤ N)
입력으로 주어지는 트리는 항상 올바른 트리임이 보장된다.
출력
Q줄에 걸쳐 각 쿼리의 답을 정수 하나로 출력한다.
풀이
루트가 정해진 트리에서 어떤 정점을 루트로 하는 서브트리 크기를 여러 번 묻는 문제다. 따라서 한 번의 DFS로 각 정점의 서브트리 크기를 전부 미리 계산해 두는 것이 가장 효율적이다.
코드에서는 루트에서 시작해 자식 정점으로 DFS를 내려가고, 자식 서브트리 크기들을 모두 더한 뒤 자기 자신 1을 더해 현재 정점의 서브트리 크기를 만든다. 이후 질의는 저장된 배열 값을 바로 출력하면 된다.
전처리 한 번으로 모든 질의를 에 답하는 구조다.
코드
import java.io.*;
import java.util.*;
public class Main {
private static ArrayList<Integer>[] list;
private static int[] dp;
public static void main(String[] args) throws IOException {
BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
StringBuilder sb = new StringBuilder();
StringTokenizer st = new StringTokenizer(br.readLine());
int N = Integer.parseInt(st.nextToken());
int R = Integer.parseInt(st.nextToken());
int Q = Integer.parseInt(st.nextToken());
dp = new int[N + 1];
list = new ArrayList[N + 1];
for(int i = 1; i < N + 1; i++) {
list[i] = new ArrayList<>();
}
Arrays.fill(dp, 1);
for (int i = 0; i < N - 1; i++) {
st = new StringTokenizer(br.readLine());
int U = Integer.parseInt(st.nextToken());
int V = Integer.parseInt(st.nextToken());
list[U].add(V);
list[V].add(U);
}
solve(R, -1);
for (int i = 0; i < Q; i++) {
int U = Integer.parseInt(br.readLine());
sb.append(dp[U]).append('\n');
}
System.out.println(sb);
}
private static void solve(int idx, int parent) {
for (int n : list[idx])
if (parent != n)
solve(n, idx);
if (parent != -1)
dp[parent] += dp[idx];
}
}복잡도
- 시간 복잡도: 전처리 , 질의당
- 공간 복잡도:
마무리
질의마다 트리를 다시 볼 필요는 없다. 서브트리 크기는 한 번의 DFS로 전부 채워 둘 수 있고, 이후에는 배열 조회만 하면 된다.
