[Algo] 백준 15681 트리와 쿼리 JAVA

문제 바로가기 : https://www.acmicpc.net/problem/15681

메모리 : 75264KB

시간 : 724ms

언어 : Java 11

---

코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.StringTokenizer;

public class Main {
    static List<List<Integer>> list;
    static int N;
    static boolean[] check;
    static int[] cnt;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        N = Integer.parseInt(st.nextToken());
        int R = Integer.parseInt(st.nextToken());
        int Q = Integer.parseInt(st.nextToken());
        list = new ArrayList<>();
        cnt = new int[N + 1];
        for (int i = 0; i <= N; i++) {
            list.add(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());
            list.get(a).add(b);
            list.get(b).add(a);
        }
        check = new boolean[N + 1];
        Arrays.fill(cnt, 1);
        check[R] = true;
        dfs(R, 0);
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < Q; i++) {
            sb.append(cnt[Integer.parseInt(br.readLine())]).append("\n");
        }
        System.out.print(sb);
    }

    private static void dfs(int node, int depth) {
        for (int i = 0; i < list.get(node).size(); i++) {
            int n = list.get(node).get(i);
            if (!check[n]) {
                check[n] = true;
                dfs(n, depth + 1);
                cnt[node] += cnt[n];
            }
        }
    }
}

풀이

인접리스트로 방향트리를 구성한 후 하위노드의 개수를 세서 저장한다.