백준 15681 트리와 쿼리 (Java)

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

1. 문제

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

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

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

입력

트리의 정점의 수 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)

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

출력

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

 

2. 풀이

 U가 주어졌을 때 본인 노드를 포함하여 자식 노드를 구하여 출력하면 된다. 다음과 같이 구할 수 있다.

• (서브 트리의 정점의 수)=(자식 서브트리의 정점의 수)+1(서브 트리의 루트 노드)

        node = new int[n + 1];

        Arrays.fill(node, 1);

 우선 노드의 수를 담을 배열 하나를 선언하고,본인 노드의 개수인 1을 담는다.

 

    static void search(int idx, int parent) {
        for (int next : g[idx]) {
            if (parent != next) {
                search(next, idx);
            }
        }

        if (parent != -1) {
            node[parent] += node[idx];
        }
    }

루트 노드에서 시작하여, 자식 노드를 찾아가며 재귀 함수를 돌려 자식 노드의 수를 부모 노드에 추가한다.

        for(int i=0; i<q; i++) {
            int idx = Integer.parseInt(reader.readLine());
            System.out.println(node[idx]);
        }

해당 연산이 끝난 후, 각 연산에 대해 자신이 부모 노드인 서브트리의 노드의 개수를 출력하면 정답을 구할 수 있다.

 

3. 코드 전문

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

public class Main {

    static int n;
    static int r;
    static int q;

    static List<Integer>[] g;
    static int[] node;

    public static void main(String[] args) throws IOException {
        BufferedReader reader = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer tokenizer = new StringTokenizer(reader.readLine());

        n = Integer.parseInt(tokenizer.nextToken());
        r = Integer.parseInt(tokenizer.nextToken());
        q = Integer.parseInt(tokenizer.nextToken());

        g = new List[n + 1];
        node = new int[n + 1];

        Arrays.fill(node, 1);

        for (int i = 0; i <= n; i++) {
            g[i] = new ArrayList<>();
        }

        for (int i = 0; i < n - 1; i++) {
            tokenizer = new StringTokenizer(reader.readLine());
            int u = Integer.parseInt(tokenizer.nextToken());
            int v = Integer.parseInt(tokenizer.nextToken());

            g[u].add(v);
            g[v].add(u);
        }


        search(r, -1);

        for (int i = 0; i < q; i++) {
            int idx = Integer.parseInt(reader.readLine());
            System.out.println(node[idx]);
        }


    }

    static void search(int idx, int parent) {
        for (int next : g[idx]) {
            if (parent != next) {
                search(next, idx);
            }
        }

        if (parent != -1) {
            node[parent] += node[idx];
        }
    }
}