[백준] 15681: 트리와 쿼리 - JAVA

30 Jul 2023 in Algorithm on 다이나믹 프로그래밍

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

풀이

트리에 속한 간선의 정보가 주어지고, 루트의 번호가 주어진다.

트리의 정점을 포함해 자신보다 하위에 있는 정점들의 개수를 구하는 문제이다.

먼저 간선의 정보를 인접리스트에 담았다.

루트의 번호부터 dfs탐색을 통해 진행하는데 정점의 개수의 최대값이 크다보니 매번 구하면 시간 초과가 발생할 것 같았다.

dp배열을 만들어 dfs탐색을 하며 dp에 저장하고 반환하는 식으로 구현했다.

---

메모리: 77800KB

시간: 812ms

언어: Java 11

코드

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

public class Main {
    static ArrayList<ArrayList<Integer>> tree;
    static int[] dp;
    static boolean[] visit;

    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());

        tree = new ArrayList<>();
        for (int i = 1; i <= N + 1; i++) {
            tree.add(new ArrayList<>());
        }

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

            tree.get(u).add(v);
            tree.get(v).add(u);
        }

        dp = new int[N + 1];
        visit = new boolean[N + 1];
        visit[R] = true;
        count(R);

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

        System.out.println(sb);
    }

    private static int count(int n) {
        int total = 1;
        for (int next : tree.get(n)) {
            if (!visit[next]) {
                visit[next] = true;
                total += count(next);
            }
        }

        return dp[n] = total;
    }

}