[알고리즘] 백준 - 트리와 쿼리

백준 - 트리와 쿼리

내 풀이

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

public class baekjoon_15681 {

    static int N, R, Q;
    static List<Integer>[] graph;
    static int[] dp;
    static boolean[] visited;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        String[] inputs = br.readLine().split(" ");
        N = Integer.parseInt(inputs[0]);
        R = Integer.parseInt(inputs[1]);
        Q = Integer.parseInt(inputs[2]);

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

        dp = new int[N + 1];
        visited = new boolean[N + 1];

        for (int i = 0; i < N - 1; i++) {
            inputs = br.readLine().split(" ");
            int A = Integer.parseInt(inputs[0]);
            int B = Integer.parseInt(inputs[1]);
            graph[A].add(B);
            graph[B].add(A);
        }

        dfs(R); //루트
        for (int i = 0; i < Q; i++) {
            int q = Integer.parseInt(br.readLine());
            System.out.println(dp[q]);
        }


    }
    //dp[root]는 root를 포함한 서브트리에 속한 정점의 개수
    private static int dfs(int root) {
        if (dp[root] != 0) {
            return dp[root];
        }
        visited[root] = true;

        int count = 1; //root 본인 한 칸

        for (int adj : graph[root]) {
            if (!visited[adj]) {
                count += dfs(adj);
            }
        }
        dp[root] = count;
        return count;
    }
}

기초적인 트리 dp문제다. 그래프를 만들고 dp 배열을 만드는데 dp[i]는 i를 루트로 하는 서브트리에서 루트를 포함한 정점의 개수라고 정의했다. 그러면 리프노드에서 1을 반환하면 된다.