[BOJ 15681] 트리와 쿼리 (Java, DP, DFS)

 

 

 

---

문제 링크

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

15681번: 트리와 쿼리

---

문제 풀이 

- 시간 복잡도 O(V + E) // 인접 리스트 사용

- 백준 링크에 들어가보면 하단에 친절하게 로직을 설명하고 있어 참고하면 됨

# 입력 트리를 그릴 경우 5번 Root 기준 
                5      
            4        6 
        3         7  8  9 
      1   2         

# subtree 결과
subtree [0, 1, 1, 3, 4, 9, 4, 1, 1, 1]

---

제출 코드

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

public class Main {
    
    static StringBuilder sb = new StringBuilder();
    static InputProcessor inputProcessor = new InputProcessor();

    static String NEW_LINE = System.lineSeparator();
    static int N, R, Q;
    static List<Integer>[] ADJ;
    static int[] SUBTREE;

    public static void main(String[] args) throws IOException {
        input();
        preprocess();

        for(int i = 1; i <= Q; i++) {
            int u = inputProcessor.nextInt();
            sb.append(SUBTREE[u]).append(NEW_LINE);
        }

        output();
    }

    private static void preprocess() {
        dfs(R, -1);
    }

    private static void dfs(int node, int prev) {
        SUBTREE[node] = 1;

        for(int child : ADJ[node]) {
            if(child == prev) continue;

            dfs(child, node);
            SUBTREE[node] += SUBTREE[child];
        }
    }

    private static void input() {
        N = inputProcessor.nextInt(); // 정점의 수
        R = inputProcessor.nextInt(); // 루트
        Q = inputProcessor.nextInt(); // 쿼리의 수

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

        for(int i = 1; i <= N - 1; i++) {
            int u = inputProcessor.nextInt();
            int v = inputProcessor.nextInt();

            ADJ[u].add(v);
            ADJ[v].add(u);
        }

        SUBTREE = new int[N + 1];
    }


    private static void output() throws IOException {
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        bw.write(sb.toString());
        bw.flush();
        bw.close();
    }

    private static class InputProcessor {
        BufferedReader br;
        StringTokenizer st;

        public InputProcessor() {
            this.br = new BufferedReader(new InputStreamReader(System.in));
        }

        public String next() {
            while(st == null || !st.hasMoreElements()) {
                try {
                    st = new StringTokenizer(br.readLine());
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }

            return st.nextToken();
        }

        public String nextLine() {
            String input = "";
            try {
                input = br.readLine();
            } catch (IOException e) {
                throw new RuntimeException(e);
            }

            return input;
        }

        public int nextInt() {
            return Integer.parseInt(next());
        }

        public long nextLong() {
            return Long.parseLong(next());
        }

    }
    
}