백준 15681 - 트리와 쿼리

문제

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

리뷰

로직을 구현하는 것은 어렵지 않았는데 계속 메모리 초과가 나서 골치를 먹었던 문제였다.
makeTree는 노드 간의 연결관계를 표현한 graph를 바탕으로 재귀적으로 트리를 구성한다.
countSubtreeNodes는 재귀적으로 nodeCount에 노드별 서브 트리 노드 수를 저장한다.

메모리를 최소화하기 위해 List<Integer>[] 형태로 graph,tree 를 정의하였다.
로직의 시간복잡도는 최대 $O(N)$이며, 이는 $N$이 $10^5$일때도 문제 시간 제한을 무난히
통과한다.

코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.List;
import java.util.StringTokenizer;

import static java.lang.Integer.parseInt;

public class Main {
    static int N, R, Q;
    static List<Integer>[] graph, tree;
    static int[] nodeCount;

    public static void main(String[] args) throws IOException {
        BufferedReader in = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(in.readLine());
        N = parseInt(st.nextToken());
        R = parseInt(st.nextToken());
        Q = parseInt(st.nextToken());

        graph = new List[N + 1];
        tree = new List[N + 1];

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

        nodeCount = new int[N + 1];

        for (int i = 0; i < N - 1; i++) {
            st = new StringTokenizer(in.readLine());
            int u = parseInt(st.nextToken());
            int v = parseInt(st.nextToken());

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

        makeTree(R, 0);
        countSubtreeNodes(R);

        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < Q; i++) {
            int node = parseInt(in.readLine());
            sb.append(nodeCount[node] + "\n");
        }

        System.out.print(sb);
        in.close();
    }

    static void makeTree(int current, int parent) {
        for (Integer node : graph[current]) {
            if (node != parent) {
                tree[current].add(node);
                makeTree(node, current);
            }
        }
    }

    static void countSubtreeNodes(int current) {
        nodeCount[current] = 1;
        for (Integer node : tree[current]) {
            countSubtreeNodes(node);
            nodeCount[current] += nodeCount[node];
        }
    }
}

결과