[백준 15681] 트리와 쿼리 / 자바 / dfs

#문제         

레벨: G5
알고리즘: dfs 
풀이시간: 
힌트 참조 유무:

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

---

#문제 풀이        

 

가장 직관적인 건 Q를 새로 입력 받을 때마다 세팅해놓은 트리의 자식 개수를 탐색하는 거다. 아래코드와 같이 말이다. 

그러나 아래와 같이 하면 시간초과가 난다.

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

public class Main {
    static List<List<Integer>> tree;

    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 = 0; i <= N; 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);
        }

        // R을 루트로 하는 트리 구조로 변경
        makeRootedTree(R, -1);

        // 쿼리 처리
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < Q; i++) {
            int U = Integer.parseInt(br.readLine());
            int subtreeSize = countSubtreeSize(U);
            sb.append(subtreeSize).append("\n");
        }

        System.out.print(sb);
    }

    // R을 루트로 하는 트리 구조로 변경
    static void makeRootedTree(int current, int parent) {
        for (int i = 0; i < tree.get(current).size(); i++) {
            int child = tree.get(current).get(i);
            if (child != parent) {
                makeRootedTree(child, current);
            } else {
                tree.get(current).remove(i);
                i--;
            }
        }
    }

    // 서브트리의 크기를 계산
    static int countSubtreeSize(int node) {
        int size = 1; // 현재 노드 포함
        for (int child : tree.get(node)) {
            size += countSubtreeSize(child);
        }
        return size;
    }
}

 

그래서 우린 다른 방법을 생각하는 것이다. 한 번의 탐색으로 node마다 자식 노드를 구하는 것이다. 

트리가 위와 같이 되어 있으면 노드5의 자식은 노드4의 자식(자기포함) + 노드6의 자식(자기포함)이고 
노드4의 자식은 노드3의 자식 수(자기포함)이다. 그러니 우린 dfs로 맨 아래로 내려갔다. 자신포함 자식을 리턴해주면 된다. 

---

#풀이 코드      

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

public class Main {
    static List<List<Integer>> tree;
    static int[] subtreeSize;

    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 = 0; i <= N; 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);
        }

        subtreeSize = new int[N + 1];
        dfs(R, -1);

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

        System.out.print(sb);
    }

    static int dfs(int node, int parent) {
        subtreeSize[node] = 1;
        for (int child : tree.get(node)) {
            if (child != parent) {
                subtreeSize[node] += dfs(child, node);
            }
        }
        return subtreeSize[node];
    }
}