BOJ. Composite Function And Query (17435)

Problem: BOJ 17435 — Composite Function And Query

한국어 · 日本語

Approach

The input defines a function f on the integers 1..M. Each query asks for f applied K times to a starting value x, i.e. f^K(x). Applying the function one step at a time can take up to 500,000 operations per query, so precompute powers of the function with binary lifting.

Let up[b][x] be the value reached after applying f exactly 2^b times to x. The base row is up[0][x] = f(x). Two consecutive jumps of length 2^(b-1) make a jump of length 2^b, so:

up[b][x] = up[b - 1][up[b - 1][x]]

For a query, inspect the set bits of K. Whenever bit b is set, replace the current value with up[b][current]. The selected jump lengths sum to K, so this computes exactly f^K(x). In particular, when K = 0 no bit is set and the answer remains x; K = 1 uses only the base row. The table has 20 rows, covering every bit of K <= 500000 (the highest possible set bit is 18), so no query can index beyond it.

Preprocessing takes O(M log Kmax) time and memory, where Kmax = 500000. Each query takes O(log Kmax) time.

import java.io.BufferedInputStream;
import java.io.IOException;

public class Main {
    private static final int LEVELS = 20;

    public static void main(String[] args) throws IOException {
        FastScanner input = new FastScanner();
        int m = input.nextInt();
        int[][] up = new int[LEVELS][m + 1];

        for (int x = 1; x <= m; x++) {
            up[0][x] = input.nextInt();
        }

        for (int bit = 1; bit < LEVELS; bit++) {
            for (int x = 1; x <= m; x++) {
                up[bit][x] = up[bit - 1][up[bit - 1][x]];
            }
        }

        int queryCount = input.nextInt();
        StringBuilder answer = new StringBuilder();
        for (int query = 0; query < queryCount; query++) {
            int k = input.nextInt();
            int value = input.nextInt();
            for (int bit = 0; bit < LEVELS; bit++) {
                if ((k & (1 << bit)) != 0) {
                    value = up[bit][value];
                }
            }
            answer.append(value).append('\n');
        }

        System.out.print(answer);
    }

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

        int nextInt() throws IOException {
            int c;
            do {
                c = read();
            } while (c <= ' ');

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