[백준/Java] 15681 : 트리와 쿼리

https://www.acmicpc.net/problem/15681

 

이번 문제는 하나의 트리에서 서브트리의 루트를 정했을 때 그 트리에 포함된 정점을 출력하는 문제입니다. 기본적인 트리 문제에 조금 더 생각을 하면 될 것 같습니다.

 

매번 서브 트리 루트가 정해질 때 마다 계산을 하기에는 정점의 수가 최대 500,000 쿼리의 수가 500,000 이고 트리 전체는 변하지 않기 때문에 dp를 통해 미리 계산을 하는 식으로 해결을 하였습니다.

 

import java.io.*;
import java.util.*;

public class Backjoon_15681_트리와_쿼리 {
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());

        int N = Integer.parseInt(st.nextToken());
        int R = Integer.parseInt(st.nextToken());
        int Q = Integer.parseInt(st.nextToken());

        boolean[] check = new boolean[N+1];
        ArrayList<Integer>[] c = new ArrayList[N+1];
        ArrayList<Integer>[] arr = new ArrayList[N+1];

        for(int i = 1 ; i <= N ; i++){
            arr[i] = new ArrayList<>();
            c[i] = new ArrayList<>();
        }

        for(int i = 0 ; i < N - 1 ; i++){
            st = new StringTokenizer(br.readLine());
            int V1 = Integer.parseInt(st.nextToken());
            int V2 = Integer.parseInt(st.nextToken());

            arr[V1].add(V2);
            arr[V2].add(V1);
        }

        Queue<Integer> q = new LinkedList<>();
        q.offer(R);
        check[R] = true;
        while(!q.isEmpty()){
            Integer poll = q.poll();
            for(int value : arr[poll]){
                if(!check[value]){
                    check[value] = true;
                    c[poll].add(value);
                    q.offer(value);
                }
            }
        }

        int[] answer = new int[N+1];
        dp(answer, c, R);

        StringBuilder sb = new StringBuilder();
        for(int i = 0; i < Q ; i++){
            int idx = Integer.parseInt(br.readLine());
            sb.append(answer[idx]).append("\n");
        }
        System.out.println(sb);
    }

    private static int dp(int[] answer, ArrayList<Integer>[] c, int r) {

        if(c[r].size() == 0){
            answer[r] = 1;
            return 1;
        }

        for(int value : c[r]){
            answer[r] += dp(answer, c, value);
        }
        answer[r]++;
        return answer[r];
    }
}