Subtask 1

Initialize an array ~A~ with the values from ~P~ and ~NXT~, where ~NXT_i~ is the position of ~i~ after applying the operation once (i.e., ~NXT_{P_i} = i~).

Next, for each of the ~K~ rounds, set ~A_i~ to ~NXT_{A_i}~ for every index ~i~ (~1 \le i \le N~).

Finally, output ~A~.

Time Complexity: ~\mathcal{O}(KN)~

Subtask 2

Redefine ~NXT_{i, k}~ as the position of ~i~ after applying the operation ~2^{k-1}~ times.

Build ~NXT_{i, k}~ for ~1 \le i \le N~ and ~1 \le k \le 61~ (since ~10^{18} \le 2^{61-1}~) by knowing that ~NXT_{i, k} = NXT_{NXT_{i, k-1}, k-1}~.

If the ~k^\text{th}~ bit is toggled in the binary representation of ~K~, set ~A_i~ to ~NXT_{A_i, k}~.

Finally, output ~A~.

Time Complexity: ~\mathcal{O}(N \log K)~

Alternative Solution

Note that binary lifting is not the only solution to this problem.

Observe that ~A~ will repeat after at most ~N~ rounds since the maximum cycle length is ~N~.

It suffices to determine the cycles and shift each by ~K~ modulo the length of the cycle.

The implementation details are left as an exercise to the reader.

Time Complexity: ~\mathcal{O}(N)~