Approach: Binary Exponentiation
This solution uses Binary Exponentiation, an approach similar to Binary Lifting. Instead of storing all powers of two in a 2D array, we calculate the next power by squaring the current mapping.
Initially, p[i] represents the destination of person i after one step. By setting p[i] = p[p[i]], we double the number of steps represented by p.
We maintain another array, res, initialized so that res[i] = i. We process the binary representation of k from right to left. Whenever the current bit is set, we apply the current mapping to res. After each iteration, we square p and shift k one bit to the right.
After processing all bits, res[i] contains the final destination of the balls originally belonging to person i. We then add a[i] to ans[res[i]] and find the maximum number of balls.
The time complexity is 𝑂(𝑛logK) , while the auxiliary space complexity is 𝑂(𝑛) . Compared with traditional Binary Lifting, this approach is more memory-efficient because it does not store every power-of-two jump.
The original editorial for this problem, which uses traditional Binary Lifting, can be found on the COCI website.