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

예제 입력 1

9 5 3
1 3
4 3
5 4
5 6
6 7
2 3
9 6
6 8
5
4
8

예제 출력 1

9
4
1

풀이

이 문제는 dfs(깊이 우선 탐색) 알고리즘과 다이나믹 프로그래밍 알고리즘을 이용해서 풀 수 있었다. dfs알고리즘을 이용해서 루트부터 트리를 구성하고 다이나믹 프로그래밍을 이용해서 서브트리의 정점의 수를 구했다.

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.ArrayList;

public class Main {
    static ArrayList<Integer>[] map;
    static ArrayList<Integer>[] tree;
    static int[] cnt;
    static boolean[] visited;

    public static void main(String[] args) throws Exception{
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        String[] input = br.readLine().split(" ");
        int N = Integer.parseInt(input[0]);
        int R = Integer.parseInt(input[1]);
        int Q = Integer.parseInt(input[2]);
        cnt = new int[N+1];
        visited = new boolean[N+1];
        tree = new ArrayList[N+1];
        map = new ArrayList[N+1];
        for(int i=1; i<=N; i++) {
            cnt[i] = 1;
            map[i] = new ArrayList<>();
            tree[i] = new ArrayList<>();
        }

        for(int i=0; i<N-1; i++) {
            input = br.readLine().split(" ");
            int start = Integer.parseInt(input[0]);
            int end = Integer.parseInt(input[1]);
            map[start].add(end);
            map[end].add(start);
        }
        
        make_tree(R, 0);
        count(R);

        StringBuilder sb = new StringBuilder();
        for(int i=0; i<Q; i++) {
            int ans = Integer.parseInt(br.readLine());
            sb.append(cnt[ans]).append("\n");
        }

        System.out.print(sb);
    }

    public static void make_tree(int temp, int parent) {

        for(int next : map[temp]) {
            if(next != parent) {
                tree[temp].add(next);
                make_tree(next, temp);
            }
        }
    }

    public static void count(int temp) {

        for(int next : tree[temp]) {
            count(next);
            cnt[temp] += cnt[next];
        }
    }
}