[백준/Java] 15681

트리와 쿼리

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

문제

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

• 정점 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줄에 걸쳐 각 쿼리의 답을 정수 하나로 출력한다.

문제 풀이 방법 및 해설

이번 문제는 하나의 트리에서 서브트리의 루트를 정했을 때 그 트리에 포함된 정점을 출력하는 문제였다. 기본적인 트리 문제에 조금 더 생각을 해야하는 문제다.

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

이전 풀이와 오늘 풀이의 차이점은 이전 풀이에는 부모를 저장하였지만 오늘 풀이는 자식들을 저장하여 dp를 한 점이 있다.

풀이 코드

이전 풀이 코드

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

public class Main {

    static int[] size;
    static int[] parent;
    static ArrayList<Integer>[] arr;
    
    public static void main(String[] args) throws IOException {

        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        StringTokenizer st = new StringTokenizer(br.readLine());

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

        arr = new ArrayList[N+1];
        parent = new int[N+1];
        size = new int[N+1];

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

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

            arr[a].add(b);
            arr[b].add(a);
        }

        Queue<Integer> q = new LinkedList();
        q.offer(R);

        while(!q.isEmpty()){
            int data = q.poll();
            for(int value : arr[data]){
                if(value == parent[data]) continue;
                q.offer(value);
                parent[value] = data;
            }
        }

        countSubTree(R);

        for(int i = 0 ; i < Q ; i++){
            int data = Integer.parseInt(br.readLine());
			System.out.println(size[data]);
        }
    }


    public static void countSubTree(int cur) {
        size[cur] = 1;

        for(int next : arr[cur]) {
            if(next == parent[cur]) continue;
            countSubTree(next);
            size[cur] += size[next];
        }
    }
}

오늘 풀이 코드

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

public class Main {

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