BOJ. Tree and Query 2 (13511)

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

Root the tree at vertex 1 and traverse it iteratively. For each vertex, record its depth and parent. The root is its own parent, so every ancestor-table entry remains valid. up[v][j] is the ancestor of v reached after 2^j edges, and weight[v][j] is the sum of those edge weights. Build both tables by doubling: the second half of a 2^j jump starts at up[v][j - 1].

For a path query, find the lowest common ancestor (LCA) by first lifting the deeper endpoint to the other endpoint’s depth, then lifting both from the largest jump down until their parents match. The path distance is the sum of the two weighted climbs from the endpoints to the LCA. The accumulator is long, since a path may contain many positive edge weights.

For a type-2 query, let upEdges be the number of edges from a to the LCA and downEdges the number from the LCA to b. The path contains upEdges + downEdges + 1 vertices. If the 1-indexed k is at most upEdges + 1, the answer is a lifted by k - 1 edges. Otherwise, count backward from b: lift it by upEdges + downEdges + 1 - k edges. This handles the LCA and both endpoints without special path manipulation.

The traversal and table construction are iterative, so a chain of 100,000 vertices cannot overflow the call stack. Preprocessing takes O(N log N) time and O(N log N) storage; each query takes O(log N) time.

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

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

        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;
        }

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

    static class Edge {
        int to;
        int weight;

        Edge(int to, int weight) {
            this.to = to;
            this.weight = weight;
        }
    }

    static int[][] up;
    static long[][] weight;
    static int[] depth;

    static int lift(int vertex, int steps) {
        for (int j = 0; steps > 0; j++, steps >>= 1) {
            if ((steps & 1) != 0) vertex = up[vertex][j];
        }
        return vertex;
    }

    static int lca(int a, int b) {
        if (depth[a] < depth[b]) {
            int temp = a;
            a = b;
            b = temp;
        }
        a = lift(a, depth[a] - depth[b]);
        if (a == b) return a;

        for (int j = up[0].length - 1; j >= 0; j--) {
            if (up[a][j] != up[b][j]) {
                a = up[a][j];
                b = up[b][j];
            }
        }
        return up[a][0];
    }

    static long climbWeight(int vertex, int steps) {
        long total = 0;
        for (int j = 0; steps > 0; j++, steps >>= 1) {
            if ((steps & 1) != 0) {
                total += weight[vertex][j];
                vertex = up[vertex][j];
            }
        }
        return total;
    }

    public static void main(String[] args) throws IOException {
        FastScanner input = new FastScanner();
        int n = input.nextInt();
        int levels = 32 - Integer.numberOfLeadingZeros(n);
        ArrayList<Edge>[] graph = new ArrayList[n];
        for (int i = 0; i < n; i++) graph[i] = new ArrayList<>();

        for (int i = 0; i < n - 1; i++) {
            int a = input.nextInt() - 1;
            int b = input.nextInt() - 1;
            int w = input.nextInt();
            graph[a].add(new Edge(b, w));
            graph[b].add(new Edge(a, w));
        }

        up = new int[n][levels];
        weight = new long[n][levels];
        depth = new int[n];
        boolean[] visited = new boolean[n];
        int[] stack = new int[n];
        int top = 0;
        stack[top++] = 0;
        visited[0] = true;
        up[0][0] = 0;

        while (top > 0) {
            int vertex = stack[--top];
            for (Edge edge : graph[vertex]) {
                if (visited[edge.to]) continue;
                visited[edge.to] = true;
                up[edge.to][0] = vertex;
                weight[edge.to][0] = edge.weight;
                depth[edge.to] = depth[vertex] + 1;
                stack[top++] = edge.to;
            }
        }

        for (int j = 1; j < levels; j++) {
            for (int vertex = 0; vertex < n; vertex++) {
                int middle = up[vertex][j - 1];
                up[vertex][j] = up[middle][j - 1];
                weight[vertex][j] = weight[vertex][j - 1] + weight[middle][j - 1];
            }
        }

        int queryCount = input.nextInt();
        StringBuilder output = new StringBuilder();
        for (int i = 0; i < queryCount; i++) {
            int type = input.nextInt();
            int a = input.nextInt() - 1;
            int b = input.nextInt() - 1;
            int ancestor = lca(a, b);

            if (type == 1) {
                output.append(climbWeight(a, depth[a] - depth[ancestor])
                        + climbWeight(b, depth[b] - depth[ancestor])).append('\n');
            } else {
                int k = input.nextInt();
                int upEdges = depth[a] - depth[ancestor];
                int downEdges = depth[b] - depth[ancestor];
                int answer = k <= upEdges + 1
                        ? lift(a, k - 1)
                        : lift(b, upEdges + downEdges + 1 - k);
                output.append(answer + 1).append('\n');
            }
        }
        System.out.print(output);
    }
}