[BOJ] 15681 : 트리와 쿼리(Java)

문제

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

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

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

입력

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

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

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

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

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

출력

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

---

💡 풀이

문제에서는 무방향으로 두 노드 u, v가 연결되어 있다는 정보만 주기때문에, 트리를 만들기 위해 루트를 기준으로 부모와 자식 관계를 설정해주는게 중요하다. 이 부분만 잘 해결한다면 어렵지 않게 풀 수 있다. 

이를 위해 처음 root에서 DFS를 수행할 때, 파라미터를 (현재 탐색 중인 노드, 그 노드의 부모 노드) 이렇게 두 개를 받아 입력된 그래프를 탐색할 때 부모 노드와 만난다면 서브트리에 속한 정점의 수를 계산하지 않고 넘어가면 된다.

 

전체적인 코드는 다음과 같다.

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Scanner;

public class Q15681 {
    static int N, R, Q;
    static ArrayList<ArrayList<Integer>> graph = new ArrayList<>();
    static int[] count;

    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);
        N = sc.nextInt();
        R = sc.nextInt(); // root
        Q = sc.nextInt();
        count = new int[N + 1];

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

        for (int i = 0; i < N - 1; i++) {
            int u = sc.nextInt();
            int v = sc.nextInt();
            graph.get(u).add(v);
            graph.get(v).add(u);
        }

        DFS(R, -1);

        for (int i = 0; i < Q; i++) {
            int q = sc.nextInt();
            System.out.println(count[q]);
        }

    }

    static void DFS(int x, int parent) {
        count[x] = 1;
        for (int y : graph.get(x)) {
            if(y == parent) continue;
            DFS(y, x);
            count[x] += count[y];
        }
    }

}