[백준] 15681. 트리와 쿼리(Java)

문제

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

 

풀이

주어진 간선을 바탕으로 무방향 그래프를 구성하고

루트부터 재귀를 돌면서 해당 노드가 루트였을 때 서비 트리의 개수를 파악하는 문제입니다.

 

재귀에서 사용하는 dp는

현재 노드가 루트였을 때의 서브트리 개수를 저장하는 배열입니다.

dp[현재노드] = 1이 기본값이며 자식 노드 탐색 후 나온 결과 값을 dp[현재노드]에 더해주는 방식으로 값이 도출됩니다.

 

이때 재귀함수의 인자로 이전 노드의 인덱스를 가져오는 이유는

무방향이기 때문에 무한루프가 도는 것을 방지하기 위한

일종의 방문처리라고 생각하면 쉽습니다.

 

코드

77,672 KB / 620 ms / 1389 B

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

public class Main {
    static int N;
    static int[] dp;
    static List<Integer>[] list;
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        StringBuilder sb = new StringBuilder();

        N = Integer.parseInt(st.nextToken());
        int R = Integer.parseInt(st.nextToken());
        int Q = Integer.parseInt(st.nextToken());

        dp = new int[N+1];
        list = new ArrayList[N+1];

        for(int i = 1 ; i <= N ; i ++)
            list[i] = new ArrayList<>();

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

            list[a].add(b);
            list[b].add(a);
        }

        recur(R, 0);

        for(int i = 0 ; i < Q ; i ++){
            int tgt = Integer.parseInt(br.readLine());
            sb.append(dp[tgt]).append("\n");
        }

        System.out.print(sb);
    }

    private static void recur(int cur, int prev) {
        dp[cur] = 1;

        for(int next: list[cur]){
            if(next == prev) continue;

            recur(next, cur);

            dp[cur] += dp[next];
        }
    }
}