[백준 P15681] 트리와 쿼리 Java

P15681 트리와 쿼리

15681번: 트리와 쿼리

문제

간선에 가중치와 방향성이 없는 임의의 루트 있는 트리가 주어졌을 때, 아래의 쿼리에 답해보도록 하자.

• 정점 U를 루트로 하는 서브트리에 속한 정점의 수를 출력한다.

만약 이 문제를 해결하는 데에 어려움이 있다면, 하단의 힌트에 첨부한 문서를 참고하자.

 

입력값

트리의 정점의 수 N과 루트의 번호 R, 쿼리의 수 Q가 주어진다. (2 ≤ N ≤ 100000, 1 ≤ R ≤ N, 1 ≤ Q ≤ 100000)

이어 N-1줄에 걸쳐, U V의 형태로 트리에 속한 간선의 정보가 주어진다. (1 ≤ U, V ≤ N, U ≠ V)

이는 U와 V를 양 끝점으로 하는 간선이 트리에 속함을 의미한다.

이어 Q줄에 걸쳐, 문제에 설명한 U가 하나씩 주어진다. (1 ≤ U ≤ N)

입력으로 주어지는 트리는 항상 올바른 트리임이 보장된다.

9 5 3
1 3
4 3
5 4
5 6
6 7
2 3
9 6
6 8
5
4
8

 

출력값

Q줄에 걸쳐 각 쿼리의 답을 정수 하나로 출력한다.

9
4
1

 

알고리즘

트리를 활용한 DP 문제이다.

입력받은 U를 시작 노드로 하여 깊이 우선 탐색 (DFS)를 진행한다.

아래의 예시를 토대로, 1번 노드를 중심으로 하는 모든 서브트리의 정점의 수를 구해보자.

 

DFS 탐색을 하며 정점 1를 루트로 하는 서브트리에 속한 정점의 수를 출력한다. 이를 위해서, 정점에서 깊이 우선 탐색을 하며 리프 노드 (5, 3, 4) 까지 이동한다. 리프 노드의 입장에서 서브트리는 자신 (1개) 뿐이므로 DP 배열에 1을 저장하고 리턴한다.

DP[3] = 1

DP[4] = 1

DP[5] = 1

정점 번호	1	2	3	4	5
서브트리의 정점 수 (DP)	0	0	1	1	1

 

함수가 순서대로 리턴된다. 자신의 자식 노드를 정점으로 하는 서브트리의 정점의 수들이 DP 배열에 저장되어 있으므로 이를 더해서 DP 배열을 업데이트 해준다. 따라서 2번 노드는 본인의 자식인 5번 노드의 DP 배열 값과 1을 더한 값을 자신의 DP 값으로 저장한다. DP[2] = DP[5] + 1

정점 번호	1	2	3	4	5
서브트리의 정점 수 (DP)	0	2	1	1	1

 

최종적으로 1번 정점까지 도달했다. 2, 3, 4번 노드가 1번 노드의 자식이므로 DP값을 더해준다. DP[1] = DP[2] + DP[3] + DP[4] + 1

정점 번호	1	2	3	4	5
서브트리의 정점 수 (DP)	5	2	1	1	1

 

이를 코드로 옮기면 아래와 같다.

 

코드

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

public class Main {
    static int N, R, Q;
    static List<Integer>[] tree;
    static boolean[] visited;
    static int[] dp;

    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());

        N = Integer.parseInt(st.nextToken());
        R = Integer.parseInt(st.nextToken());
        Q = Integer.parseInt(st.nextToken());

        dp = new int[N + 1];
        tree = new List[N + 1];
        visited = new boolean[N + 1];
        for (int i = 1; i <= N; i++) {
            tree[i] = new ArrayList<>();
        }

        for (int i = 1; i < N; i++) {
            st = new StringTokenizer(br.readLine());
            int from = Integer.parseInt(st.nextToken());
            int to = Integer.parseInt(st.nextToken());

            tree[from].add(to);
            tree[to].add(from);
        }

        findSubTree(R);

        for (int i = 0; i < Q; i++) {
            int node = Integer.parseInt(br.readLine());
            System.out.println(dp[node]);
        }
    }

    public static int findSubTree(int root) {
        visited[root] = true;
        if (dp[root] != 0) return dp[root];

        int total = 0;
        for (Integer child : tree[root]) {
            if (!visited[child])
                total += findSubTree(child);
        }

        return dp[root] = total + 1;
    }
}