[ BOJ ][JAVA][15681] 트리와 쿼리

www.acmicpc.net/problem/15681

15681번: 트리와 쿼리

시간 제한	메모리 제한	제출	정답	맞은 사람	정답 비율
1 초	128 MB	1969	1068	799	53.914%

문제

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

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

예제 입력 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

코드

/*
    트리와 쿼리.
    tree + DP
*/
import java.io.*;
import java.util.*;

public class p15681 {
    static int N, R, Q;
    static List<Integer>[] list;
    static int[] dp;
    public static void main(String[] args) throws IOException {
        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];
        list = new ArrayList[N + 1];
        for (int i = 0; i <= N; i++) {
            list[i] = new ArrayList<>();
        }

        // 트리의 간선 셋팅
        for (int i = 0; i < N - 1; i++) {
            st = new StringTokenizer(br.readLine());
            int u = Integer.parseInt(st.nextToken());
            int v = Integer.parseInt(st.nextToken());

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

        // 루트의 번호 : R, 루트의 부모는 없으니 -1
        makeTree(R, -1);

        StringBuilder sb = new StringBuilder();
        while (Q-- > 0) {
            int U = Integer.parseInt(br.readLine());
            // U를 root로 하는 서브트리에 속한 정점의 수는? -> U의 자식의 수는?
            sb.append(dp[U]).append("\n");
        }
        System.out.println(sb.toString());
    }

    // list를 바탕으로 SubTree를 만든다.
    // 루트노드는 부모가 없으므로 p가 -1이다.
    static int makeTree(int cur, int p) {
        // 이미 방문한 곳이라면 그대로 리턴
        if (dp[cur] != 0) {
            return dp[cur];
        }
        dp[cur] = 1;

        for (int child : list[cur]) {
            // 현재 노드의 붙어있는 노드가 부모노드가 아니라면 전부다 서브트리에 저장한다.
            if (child != p) {
                // 재귀 수행
                dp[cur] += makeTree(child, cur);
            }
        }

        return dp[cur];
    }

}

다른풀이

public class Main {
    public static LinkedList<Integer>[] adjList;
    public static LinkedList<Integer>[] subList;
    public static int[] parent;
    public static int[] size;
    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());
        adjList = new LinkedList[N+1];
        subList = new LinkedList[N+1];
        parent = new int[N+1];
        size = new int[N+1];
        for (int i = 1; i <= N; i++) {
            adjList[i] = new LinkedList<Integer>();
            subList[i] = new LinkedList<Integer>();
        }
        for(int i = 0; i < N-1; i++) {
            st = new StringTokenizer(br.readLine());
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());
            adjList[a].add(b);
            adjList[b].add(a);
        }

        makeTree(R, -1);
        countSubTree(R);
        for(int i = 0; i < Q; i++) {
            int q = Integer.parseInt(br.readLine());
            System.out.println(size[q]);
        }
        br.close();
    }
    public static void makeTree(int cur, int p) {
        for(int next : adjList[cur]) {
            if(next != p) {
                subList[cur].add(next);
                parent[next] = cur;
                makeTree(next, cur);
            }
        }
    }
    public static void countSubTree(int cur) {
        size[cur] = 1;
        for(int next : subList[cur]) {
            countSubTree(next);
            size[cur] += size[next];    
        }
    }
}