Q. 15681 자바 : 트리와 쿼리

Q. 15681 자바 : 트리와 쿼리

 

문제 : boj.kr/15681

골드 5 난이도의 문제이다.

 

트리와 간선이 주어진다.

이후 Q개의 정점이 주어지고 각 정점을 루트로 하는 서브트리의 정점의 수를 구하는 문제이다.

 

아래와 같은 풀이로 풀었다.

0. 간선이 주어지므로 ArrayList를 사용해 트리 구조를 만든다.

1. 루트 노드가 주어지므로, dfs 탐색을 이용해 탐색한다.

2. 서브트리의 정점의 수를 구하기 위해 탐색 이후 종료 전에

자식을 루트로 하는 서브트리의 정점의 수를 저장한다.(메모제이션 방식)

3. 이후 Q개의 정점에 대해 각 정점을 루트로 하는 서브트리의 정점의 수를 2번에서 저장한 값을 이용해 결과를 출력한다.

 

시간 복잡도는 정점의 수 N 이 최대 10,000이고 탐색하는 정점 Q도 최대 10,000이다.

그래서 아래와 같이 계산할 수 있다.

1. 서브트리 정점의 수 계산을 위한 dfs - O(log N)

2. Q개의 정점 탐색 - O(Q)

따라서 시간복잡도는 약 O(Q log N)이므로 시간초과는 걱정하지 않아도 된다.

 

아래와 같이 풀었다.

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

public class Main {
    static FastReader scan = new FastReader();
    static StringBuilder sb = new StringBuilder();

    static int n, R, Q;
    static ArrayList<Integer>[] tree;

    static int[] Dis;

    static void input() throws IOException {
        n = scan.nextInt();
        R = scan.nextInt();
        Q = scan.nextInt();
        tree = new ArrayList[n + 1];
        Dis = new int[n + 1];
        for (int i = 1; i <= n; i++) tree[i] = new ArrayList<>();
        for (int i = 1; i < n; i++) {
            int v1 = scan.nextInt();
            int v2 = scan.nextInt();
            tree[v1].add(v2);
            tree[v2].add(v1);
        }
    }

    static void dfs(int x, int prev) {

        Dis[x] = 1;
        for (int t: tree[x]) {
            if (t == prev) continue;
            dfs(t, x);
            Dis[x] += Dis[t];
        }
    }

    static void pro() {

        dfs(R, -1);

        while (Q-- > 0) sb.append(Dis[scan.nextInt()]).append('\n');
        System.out.println(sb);
    }


    public static void main(String[] args) throws IOException {
        input();
        pro();
    }


    static class FastReader {
        BufferedReader br;
        StringTokenizer st;

        public FastReader() {
            br = new BufferedReader(new InputStreamReader(System.in));
        }

        public FastReader(String s) throws FileNotFoundException {
            br = new BufferedReader(new FileReader(new File(s)));
        }

        String next() {
            while (st == null || !st.hasMoreElements()) {
                try {
                    st = new StringTokenizer(br.readLine());
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
            return st.nextToken();
        }

        int nextInt() {
            return Integer.parseInt(next());
        }

        long nextLong() {
            return Long.parseLong(next());
        }

        double nextDouble() {
            return Double.parseDouble(next());
        }

        String nextLine() {
            String str = "";
            try {
                str = br.readLine();
            } catch (IOException e) {
                e.printStackTrace();
            }
            return str;
        }
    }
}