[백준 15681] 트리와 쿼리 (트리 DP)

15681번: 트리와 쿼리 (acmicpc.net)

15681번: 트리와 쿼리

 

트리에서 DP 를 사용하는 입문 문제!

 

문제는 단순하다.

 

Q 개의 쿼리가 들어오면 이에 따른 출력을 해주면 된다.

 

쿼리는 노드 V 에 대한 서브 트리를 구성하는 노드들의 개수를 반환하는 것이다.

 

입력 조건을 보면,

 

트리의 크기도 상당히 큰 편이기도 하지만,

 

쿼리가 최대 100,000 개가 들어올 수 있다.

 

노드 V 에 대한 서브 트리를 구성하는 노드 개수를 카운팅하는 가장 쉬운 방법인

 

V 에 대한 dfs 탐색을 쿼리마다 수행한다면

 

사실, 계산해보지 않아도 시간 초과가 발생할 것임을 알 수 있다.

 

시간 초과를 방지하기 위해 반복된 연산을 메모이제이션하여 dp 테이블에 저장하고,

 

쿼리가 발생하면 그냥 dp 에서 꺼내서 출력하도록 하자.

 

아래와 같은 트리가 주어진다면,

 

출처: geeks for geeks

1 번 노드에 대한 서브 노드를 구성하는 정점은 개수는 자기 자신을 포함한 7

2 번 노드에 대한 서브 노드를 구성하는 정점은 개수는 자기 자신을 포함한 3

3 번 노드에 대한  서브 노드를 구성하는 정점은 개수는 자기 자신을 포함한 3

4,5,6,7 번 노드에 대한  서브 노드를 구성하는 정점은 개수는 자기 자신을 포함한 1

 

생각해보면

 

1번 노드에 대한 서브 노드를 구성하는 정점의 개수는

  =  2 번 노드에 대한  서브 노드를 구성하는 정점은 개수 + 3 번 노드에 대한  서브 노드를 구성하는 정점은 개수

이다.

 

2 번 노드에 대한  서브 노드를 구성하는 정점은 개수는

  =  4 번 노드에 대한  서브 노드를 구성하는 정점은 개수 + 5 번 노드에 대한  서브 노드를 구성하는 정점은 개수이다.

 

이와 같이 큰 문제가 작은 문제로 쪼개지는 특성을 이용하여 트리에 대한 서브 노드 개수를 모두 계산하는 함수를 작성할 수 있다.

 

dp 테이블을 0 으로 초기화하고,

 

root 노드에 방문한다. (여기서는 1)

 

그리고 서브 노드 개수를 카운팅할 변수인 ret 을 1 로 초기화한다.

 

그리고 dp[1] 을 확인했을 때 0 이라면, 더 작은 문제를 확인한다.(연결 노드를 탐색)

 

1과 연결된 2, 3 노드에 대하여 dp[2] , dp[3] 이 0 이라면 2 번, 3 번 노드에 대해서 더 작은 문제를 확인하고

 

ret 에 그 결과값을 누적시킨다.

 

그리고 나서 dp[현재 탐색하고 있는 노드] 에 ret 을 차례대로 저장하다보면 결국 전체 트리의 dp 테이블이 채워지는 구조이다.

 

사실 말로 설명하면 잘 와닿지 않을 것 같다..

 

아래 코드를 보고 이해하는 게 더 쉬울 것이다.

 

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

public class Main {

    /*
    트리의 정점의 수 N과 루트의 번호 R, 쿼리의 수 Q가 주어진다. (2 ≤ N ≤ 105, 1 ≤ R ≤ N, 1 ≤ Q ≤ 105)
    이어 N-1줄에 걸쳐, U V의 형태로 트리에 속한 간선의 정보가 주어진다. (1 ≤ U, V ≤ N, U ≠ V)
    이는 U와 V를 양 끝점으로 하는 간선이 트리에 속함을 의미한다.
    이어 Q줄에 걸쳐, 문제에 설명한 U가 하나씩 주어진다. (1 ≤ U ≤ N)
    입력으로 주어지는 트리는 항상 올바른 트리임이 보장된다.
     */
    static int n;
    static int r;
    static int q;
    static ArrayList<Integer>[] tree;
    static boolean[] visited;
    static int[] dp;

    public static void main(String[] args) throws IOException {
        System.setIn(new FileInputStream("src/input.txt"));
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        StringTokenizer stk;
        stk = new StringTokenizer(br.readLine());
        int n = Integer.parseInt(stk.nextToken());
        int r = Integer.parseInt(stk.nextToken());
        int q = Integer.parseInt(stk.nextToken());

        tree = new ArrayList[n + 1];
        for (int i = 0; i <= n; i++) {
            tree[i] = new ArrayList<>();
        }
        for (int i = 0; i < n - 1; i++) {
            stk = new StringTokenizer(br.readLine());
            int a = Integer.parseInt(stk.nextToken());
            int b = Integer.parseInt(stk.nextToken());
            tree[a].add(b);
            tree[b].add(a);
        }
        dp = new int[n + 1];
        visited = new boolean[n + 1];

        find(r);

        while (q-- > 0) {
            int t = Integer.parseInt(br.readLine());
            bw.write(dp[t] + "\n");
        }
        bw.flush();
        bw.close();
    }

    static int find(int v) {
        visited[v] = true;
        int ret = 1;
        if (dp[v] == 0) {
            for (int iter : tree[v]) {
                if (!visited[iter]) {
                    ret += find(iter);
                    visited[iter] = true;
                }
            }
            dp[v] = ret;
        }
        return dp[v];
    }
}