백준 - 트리와 쿼리 (JAVA)

1. 문제

 

 

2. 풀이 방법

문제를 이해하기 위해 예제로 주어진 트리를 그려봤다. 루트가 주어지기에 자연스럽게 계층 구조가 형성된다. 서브트리의 개수는 곧 자기와 자식들의 정점 수를 모두 합친 수였다.

 

시간복잡도를 따져봤을 때 정점의 개수는 십만이고, 쿼리의 개수도 십만이기에 O(N)이나 O(logN)으로 풀어야할 것 같았다.

어떻게하면 한 번의 쿼리 때 logN의 시간복잡도를 가질 수 있을까 고민해보다가 그냥 한 번 전체탐색을 돌려서 모든 정점의 서브트리의 개수를 구한 후 쿼리에 맞춰 답을 출력해주면 될 것 같았다.

 

서브트리를 구하는 방법은 재귀적으로 자식 노드들을 방문하면서 개수를 구하면 될 것 같아서 바로 코드로 옮겨봤다.

우선 트리를 형성하기 위해 ArrayList 배열을 이용하였다. 아래와 같이 주어진 입력태로 트리 자료구조를 형성한다.

		tree = new ArrayList[N+1];
        for(int i=0; i<tree.length; i++){
            tree[i] = new ArrayList<Integer>();
        }

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

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

 

 

• 서브트리 개수 구하기

주어진 루트를 기준으로 `getSubTreeN(R)`을 호출하면 탐색 시작이다. 탐색이 모두 완료되면 `subTreeN` 배열에는 각 정점에 대한 서브트리의 수가 저장될 것이다.

 

정점 `n`에 대해 자식들을 루프 돌면서 서브트리의 개수를 구한 후 다 더해준다. 그리고 마지막으로 자기 정점도 포함해야하니 1을 더해주고 리턴한다. 중요한 점은 방문 여부를 체크하는 `visited` 배열을 사용해야한다는 것이다. 입력으로 주어진 자료는 계층 관계가 없기 때문에 부모 노드에서 자식노드를 호출하였을 때, 자식 노드는 부모 노드와 연결되어 있기 때문에 다시 부모 노드를 호출한다. 그렇기 때문에 부모 노드를 방문했다고 기록을 해야해서 나는 그 방법으로 visited 배열을 이용했다. 

 public int getSubTreeN(int n){
        isVisited[n] = true;
        // 말단 노드
        if(tree[n].size() == 1 && n != R){
            return (subTreeN[n] = 1);
        }

        // 자식 노드의 서브트리 개수 알아내기
        for(int child : tree[n]){
            if(isVisited[child]){
                continue;
            }
            subTreeN[n] += getSubTreeN(child);
        }
        isVisited[n] = false;

        // 현재 노드를 개수에 포함
        return subTreeN[n] = subTreeN[n] + 1;
    }

 

3. 전체 코드

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.util.ArrayList;
import java.util.LinkedList;
import java.util.Queue;
import java.util.StringTokenizer;


class Main {

    public static BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
    public static BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));

    ArrayList<Integer> [] tree;
    public static int N;
    public static int R;
    public static int Q;
    int [] subTreeN;
    boolean [] isVisited;

    public void run () throws IOException {
        StringTokenizer st = new StringTokenizer(br.readLine());
        N = Integer.parseInt(st.nextToken());   // 트리 정점의 수
        R = Integer.parseInt(st.nextToken());   // 루트 번호
        Q = Integer.parseInt(st.nextToken());   // 쿼리의 수

        tree = new ArrayList[N+1];
        for(int i=0; i<tree.length; i++){
            tree[i] = new ArrayList<Integer>();
        }

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

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

        subTreeN = new int[N+1];
        isVisited = new boolean[N+1];

        getSubTreeN(R);

        // 쿼리 실행
        for(int i=0; i<Q; i++){
            int n = Integer.parseInt(br.readLine());
            bw.write(subTreeN[n]+"\n");
        }
        bw.flush();
    }

    public int getSubTreeN(int n){
        isVisited[n] = true;
        // 말단 노드
        if(tree[n].size() == 1 && n != R){
            return (subTreeN[n] = 1);
        }

        // 자식 노드의 서브트리 개수 알아내기
        for(int child : tree[n]){
            if(isVisited[child]){
                continue;
            }
            subTreeN[n] += getSubTreeN(child);
        }
        isVisited[n] = false;

        // 현재 노드를 개수에 포함
        return subTreeN[n] = subTreeN[n] + 1;
    }

    public static void main(String[] args) throws IOException {

        new Main().run();
    }
}