Official

E - Permute K times 2 Editorial by MMNMM


頂点 \(1,\) 頂点 \(2,\ldots,\) 頂点 \(N\) からなり、\(i\to P _ i\ (1\leq i\leq N)\) の \(N\) 辺がある有向グラフを考えます。
このグラフにおいて、\(i\) から \(k\) 回進んだ頂点を \(P ^ k _ i\) と呼ぶことにします(どの頂点についてもその頂点から出る辺は \(1\) 本のみであるため、これは一意に定まります)。

求める答えは \((P ^ {2 ^ K} _ i) _ i\) です(証明は省略しますが、\(K\) についての帰納法で示すのが平易だと思います)。

有向グラフはいくつかのサイクルに分けることができます。 それぞれのサイクルについて、\(2 ^ K\) をそのサイクルの長さで割った余りを考えれば最終的な値を求めることができます。

サイクルの個数を \(M\) として時間計算量は \(O(N+M\log K)\) となります。

実装例は以下のようになります。

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

    // 求める順列
    vector<unsigned> ans(N);
    // used[i] := i に対する答えをすでに求めていれば true
    basic_string used(N, false);
    for (unsigned i{}; i < N; ++i) if (!used[i]) {
        // i を含むサイクルを求める
        vector<unsigned> cycle;
        {
            unsigned j{i};
            while (!used[j]) {
                used[j] = true;
                cycle.push_back(j);
                j = P[j];
            }
        }
        // サイクルに対して答えを求める
        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: