E - Permute K times 2 解説 by MMNMM

より高速な解法

全体で \(O(N)\) 時間で解く方法について解説します。

各サイクルについて、その長さを \(m\) として \(O(m)\) 時間で答えを求められればよいです。

これは次の \(2\) つの方針で実現できます。

  • 事前に \(O(N)\) 時間で \(\varphi(1),\varphi(2),\ldots,\varphi(N)\) を求めておき、\(2 ^ K\bmod m\) を \(O(\log\varphi(m))\) 時間で求める
  • \(2 ^ k\bmod m\) がたかだか \(m+1\) 回で重複するため、周期性を利用して \(2 ^ K\bmod m\) を \(O(m)\) 時間で求める

どちらの方針でも全体で \(O(N)\) 時間となります。

実装例は以下のようになります。 この実装例では、ひとつめの方針で実装を行っています。

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

int main() {
    using namespace std;
    unsigned N;
    unsigned long K;
    cin >> N >> K;

    // φ(1), φ(2), ..., φ(N) を O(N) 時間で求める
    vector<unsigned> sieve(N + 1, 1), least_prime_factor(N + 1), primes, totient(N + 1, 1);
    for (unsigned i{2}; i <= N; ++i) {
        if (sieve[i]) {
            primes.push_back(i);
            least_prime_factor[i] = i;
            totient[i] = i - 1;
        }
        for (const auto &p : primes) {
            if (i * p > N)
                break;
            sieve[i * p] = false;
            least_prime_factor[i * p] = p;
            if (least_prime_factor[i] == p) {
                totient[i * p] = totient[i] * p;
                break;
            }
            totient[i * p] = totient[i] * (p - 1);
        }
    }

    vector<unsigned> P(N);
    for (auto &p : P) {
        cin >> p;
        --p;
    }

    // 2^K mod m を O(m) 時間で求める
    const auto safe_fast_pow{[&totient](unsigned a, unsigned long K, unsigned m) {
        return atcoder::pow_mod(a, min(K, K % totient[m] + totient[m]), m);
    }};

    vector<unsigned> ans(N), used(N);
    for (unsigned i{}; i < N; ++i) if (!used[i]) {
        vector<unsigned> cycle;
        {
            unsigned j{i};
            while (!used[j]) {
                used[j] = 1;
                cycle.push_back(j);
                j = P[j];
            }
        }
        const auto cycle_length{size(cycle)};
        const auto shift{safe_fast_pow(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;
}

投稿日時:
最終更新: