ALGORITHM NOTE1

BOJ 15681 - 트리와 쿼리

트리? 쿼리?

#algorithm#boj#gold#dp#graph-theory#graph-traversal#tree#dfs#tree-dp
아카이브로 돌아가기

문제 링크

문제

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

  • 정점 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을 더해 현재 정점의 서브트리 크기를 만든다. 이후 질의는 저장된 배열 값을 바로 출력하면 된다.

전처리 한 번으로 모든 질의를 O(1)O(1)에 답하는 구조다.

코드

java
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];			
	}
	
}

복잡도

  • 시간 복잡도: 전처리 O(N)O(N), 질의당 O(1)O(1)
  • 공간 복잡도: O(N)O(N)

마무리

질의마다 트리를 다시 볼 필요는 없다. 서브트리 크기는 한 번의 DFS로 전부 채워 둘 수 있고, 이후에는 배열 조회만 하면 된다.