Official

E - Permute K times 2 Editorial by en_translator


Consider a directed graph with vertices \(1,\), \(2,\ldots,\) and \(N\), and with \(N\) edges \(i\to P _ i\ (1\leq i\leq N)\).
In this graph, let us denote by \(P ^ k _ i\) the vertex you reach by advancing from \(i\), \(k\) times. (Since each vertex has a unique outgoing edge, the destination is uniquely determined.

The sought answer is \((P ^ {2 ^ K} _ i) _ i\). (We omit the proof, but we think it is straightforward to prove by induction on \(K\).)

The directed graph can be decomposed into cycles. The answer can be obtained by finding the remainder when dividing \(2 ^ K\) by the length of the cycle.

The complexity is \(O(N+M\log K)\), where \(M\) is the number of cycles.

The sample code is as follows.

#include <atcoder/math>
#include <iostream>
#include <vector>

int main() {
    using namespace std;
    unsigned N;
    unsigned long K;
    cin >> N >> K;
    vector<unsigned> P(N);
    for (auto &p : P) {
        cin >> p;
        --p;
    }

    // The sought permutation
    vector<unsigned> ans(N);
    // used[i] := true if the answer for i is already found
    basic_string used(N, false);
    for (unsigned i{}; i < N; ++i) if (!used[i]) {
        // find the cycle containing i
        vector<unsigned> cycle;
        {
            unsigned j{i};
            while (!used[j]) {
                used[j] = true;
                cycle.push_back(j);
                j = P[j];
            }
        }
        // find the cycle against this cycle
        const auto cycle_length{size(cycle)};
        const auto shift{atcoder::pow_mod(2, K, cycle_length)};
        for (unsigned j{}; j < cycle_length; ++j)
            ans[cycle[j]] = cycle[(j + shift) % cycle_length];
    }
    
    for (const auto p : ans)
        cout << p + 1 << " ";
    cout << endl;
    return 0;
}

posted:
last update: