BOJ 15681 트리와 쿼리

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

문제

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

• 정점 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줄에 걸쳐 각 쿼리의 답을 정수 하나로 출력한다.

---

예제 입력 1

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

예제 출력 1

9
4
1

---

코드

처음 읽을 때는 문제가 조금 어렵다고 느껴졌는데, 사실은 되게 간단하다.

트리의 간선들과 루트 노드를 가지고 트리를 만들어, 각 노드의 child 노드 수를 세어주면 된다.

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.List;
import java.util.StringTokenizer;

public class Main {

    static int[] nodeCount;
    static List<Integer>[] linked;

    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());

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

        for (int i = 1; i < N; i++) {
            st = new StringTokenizer(br.readLine());
            int nodeA = Integer.parseInt(st.nextToken());
            int nodeB = Integer.parseInt(st.nextToken());
            linked[nodeA].add(nodeB);
            linked[nodeB].add(nodeA);
        }

        countChildNodes(R, 0);

        // 각 쿼리의 정답 정점 수를 출력
        StringBuilder answer = new StringBuilder();
        for (int i = 0; i < Q; i++) {
            int U = Integer.parseInt(br.readLine());
            answer.append(nodeCount[U]).append('\n');
        }

        System.out.println(answer);
    }

    static int countChildNodes(int node, int parent) {
        int count = 1; // 서브 트리에 속한 정점은 검사 대상 노드도 포함하기 때문에 1로 초기화

        for (int check: linked[node]) {
            if(check != parent) {
                count += countChildNodes(check, node);
            }
        }

        return nodeCount[node] = count;
    }
}