[백준] 15681번 - 트리와 쿼리

문제: 15681번 트리와 쿼리

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

15681번: 트리와 쿼리

문제 요약

R이라는 루트를 가진 트리에서 정점 Q1, Q2 ... Qn에 대한 서브트리의 노드 갯수를 구하라.

 

예제

노드의 수 N = 9 , 루트의 번호 R = 5, 쿼리의 수 Q = 3

아래 8줄은 간선 정보 (1,3) , (4,3) , (5,4) .... (6,8)

아래 3줄은 쿼리 (Qn번 노드의 서브트리 노드 개수)

 

문제 유형

트리 동적 계획법

 

문제 풀이

1. 서브트리의 노드 수를 subtree_size 배열에 저장한다. 

2. 깊이우선탐색을 이용해 자식과 부모를 설정하며 재귀 호출을 반복한다. 

3. 재귀에서 빠져나왔다는 것은 자신과 연결되어있는 자식 노드의 처리가 끝났다는 뜻이므로, 자기자신의 노드 사이즈인 1을 subtree_size에 설정하고, 자신의 자식 노드의 subtree_size를 가산한다. 

 

문제 고려사항

1. root 노드로부터 dfs를 수행할 때, 부모 노드는 없기 때문에 -1로 넘긴다. 

2. 노드 v와 연관된 노드를 순서대로 찾아 수행할 때, 노드 v의 간선 정보 리스트에는 부모 노드도 저장되어 있다. 탐색은 자식으로 이어져야 하지 부모로 역류하면 안되기 때문에, 부모 노드에 대한 탐색은 무시한다.

3. 최대 N이 10의 5승이기 때문에 단순히 N*N 배열을 만들면 메모리 초과 문제가 발생한다. 이번에는 ArrayList배열을 생성하여 연결되는 간선들만 list에 저장한다. (배열을 사용할 때면 연결되지 않는 간선들의 인덱스 값은 항상 더미로 채워져야 하기 때문에 메모리 낭비가 심하다.)

4. Q가 10의 5승이기 때문에 출력 부담이 심할 것이라 생각하여 stringBuffer 을 사용한다.

 

전체코드

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,R,Q,U;
    static List<Integer>[] tree;
    static int subtree_size[];
    private static void dfs(int v, int p) {
        for(int c : tree[v]){
            if(c==p) continue; //부모 노드로 역류 금지
            dfs(c,v); //다음 노드를 자식으로, 현재 노드를 부모 노드로
        }
        subtree_size[v] = 1;
        for(int c: tree[v]){ //v와 간선이 있는 c를 찾는다.
            if(c==p) continue; //부모 노드는 제외. 자신의 서브트리 노드 크기를 구하는 것이기 때문이다.
            subtree_size[v] += subtree_size[c]; //자식노드의 서브트리 크기를 찾아서 자신의 서브트리 크기에 가산한다.
        }
    }
    public static void main(String[] args) throws IOException {
        BufferedReader bf = new BufferedReader(new InputStreamReader(System.in));
        StringBuffer sb = new StringBuffer();
        StringTokenizer st = new StringTokenizer(bf.readLine());

        N = Integer.parseInt(st.nextToken());
        R = Integer.parseInt(st.nextToken());
        Q = Integer.parseInt(st.nextToken());
        tree = new ArrayList[N+1];
        for(int i=1; i<N+1; i++) {
            tree[i] = new ArrayList<>();
        }
        subtree_size = new int[N+1];
        for (int i = 0; i < N - 1; i++) {
            st = new StringTokenizer(bf.readLine());
            int n = Integer.parseInt(st.nextToken());
            int m = Integer.parseInt(st.nextToken());
            tree[n].add(m); //tree에 간선정보 등록
            tree[m].add(n); //무향 그래프이기 때문에 서로서로 간선 정보를 등록 
        }
        dfs(R,-1);
        for(int i=0; i<Q; i++){
            int inp = Integer.parseInt(bf.readLine());
            sb.append(subtree_size[inp]).append("\n");
        }
        System.out.println(sb.toString());
    }
}