BOJ. Tree And Query (15681)

Problem: BOJ 15681 — Tree and Query · 한국어 · 日本語

Iterative rooting and reverse-order accumulation

Root the undirected tree at R. During a stack-based traversal, record each node’s parent and append the node to an order array. A node is added to the stack only from its parent, so every node is pushed and recorded exactly once; when an edge leads back to the current node’s parent, skip it. Because a parent is processed before its children are discovered, every parent appears before its children in order.

Initialize every subtree size to one, counting the node itself. Then scan order backward and add each node’s accumulated size to its parent. Reversing the parent-before-child order guarantees that all children have contributed before their parent is propagated upward. Thus, when preprocessing finishes, subtreeSize[u] is exactly the number of nodes in the subtree rooted at u. A query is answered directly from this array in O(1) time; this includes the root query (which returns N) and a leaf query (which returns 1).

The traversal and accumulation each visit every node once, and reading the N - 1 edges takes O(N) time. Total preprocessing is O(N), followed by O(Q) time for queries. The adjacency lists, parent/order arrays, stack, and subtree-size array use O(N) memory. No recursive calls are used, so a chain of 100,000 nodes does not consume the Java call stack.

import java.io.BufferedInputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

public class Main {
    public static void main(String[] args) throws IOException {
        FastScanner input = new FastScanner();
        int nodeCount = input.nextInt();
        int root = input.nextInt();
        int queryCount = input.nextInt();

        List<Integer>[] graph = new List[nodeCount + 1];
        for (int node = 1; node <= nodeCount; node++) {
            graph[node] = new ArrayList<>();
        }
        for (int i = 0; i < nodeCount - 1; i++) {
            int u = input.nextInt();
            int v = input.nextInt();
            graph[u].add(v);
            graph[v].add(u);
        }

        int[] parent = new int[nodeCount + 1];
        int[] order = new int[nodeCount];
        int[] stack = new int[nodeCount];
        int orderSize = 0;
        int top = 0;
        stack[top] = root;
        parent[root] = -1;

        while (top >= 0) {
            int node = stack[top--];
            order[orderSize++] = node;
            for (int neighbor : graph[node]) {
                if (neighbor == parent[node]) {
                    continue;
                }
                parent[neighbor] = node;
                stack[++top] = neighbor;
            }
        }

        int[] subtreeSize = new int[nodeCount + 1];
        for (int node = 1; node <= nodeCount; node++) {
            subtreeSize[node] = 1;
        }
        for (int i = orderSize - 1; i > 0; i--) {
            int node = order[i];
            subtreeSize[parent[node]] += subtreeSize[node];
        }

        StringBuilder output = new StringBuilder();
        for (int i = 0; i < queryCount; i++) {
            output.append(subtreeSize[input.nextInt()]).append('\n');
        }
        System.out.print(output);
    }

    private static class FastScanner {
        private final BufferedInputStream input = new BufferedInputStream(System.in);
        private final byte[] buffer = new byte[1 << 16];
        private int position;
        private int length;

        private int read() throws IOException {
            if (position == length) {
                length = input.read(buffer);
                position = 0;
                if (length == -1) {
                    return -1;
                }
            }
            return buffer[position++];
        }

        private int nextInt() throws IOException {
            int c;
            do {
                c = read();
            } while (c <= ' ' && c != -1);

            int value = 0;
            while (c > ' ') {
                value = value * 10 + c - '0';
                c = read();
            }
            return value;
        }
    }
}