2020牛客暑期多校训练营（第二场）J - Just Shuffle

题目

Given a permutation with size \(n\) and an integer \(k\), you should find a permutation substitution \(P\) that \(1, 2, \cdots, n\) will become \(A\) after performing substitution \(P\) for exactly \(k\) times. Print the permutation after performing \(P\) for once on \(1, 2, \cdots, n\). If there are multiple solutions, print any of them. If there is no solution, print "-1" in one line.

输入

The first line contains two integers \(n, k\left(1 \leq n \leq 10^{5}, 10^{8} \leq k \leq 10^{9}\right)\)
The second line contains \(n\) integers \(A_{1}, A_{2}, \cdots, A_{n}\), denoting the permutation \(A\).
It is guaranteed that \(k\) is a prime number.

输出

If there exists solutions, print n integers in one line, or print "-1" in one line.

样例输入

3 998244353
2 3 1

样例输出

3 1 2

题意

给定大小为 \(n\) 的置换 \(B\) 和大质数 \(k\) ，找到一个置换 \(A\) 使得 \(A^k = B\)。如果不存在这样的置换 \(B\) 输出 -1。

题解

对于式子 \(A^k = B\) 两边同时取 \(t\) 次方，得到：

\[A^{kt} = B^t \]

假若置换 \(A\) 的循环节长度为 \(m\) ，那么显然，\(B\) 的循环节长度也为 \(m\) 。假若我们能找到这样一个 \(t\) ，使得

\[kt \equiv 1 \mod m \]

那么求出 \(B^t\) 就能得到 \(A\) 了。显然，\(t\) 是 \(k\) 的逆元，由于 \(k\) 是大质数，所以逆元一定存在。用扩展欧几里得算一下就行了。\(k\) 是大质数保证了这样是一定能解的。

代码

#include <bits/stdc++.h>
#define lowbit(x) ((x)&(-x))
#define mem(i, a) memset(i, a, sizeof(i))
#define sqr(x) ((x)*(x))
#define ls(x) (x << 1)
#define rs(x) (x << 1 | 1)
typedef long long ll;
const double eps = 1e-8;
const double pi = acos(-1.0);
const int inf = 0x3f3f3f3f;
const int maxn = 1e5 + 7;
using namespace std;
int arr[maxn], ans[maxn], n, k;
bool vis[maxn];
void exgcd(ll a, ll b, ll &x, ll &y){
    if(!b)
        x = 1, y = 0;
    else{
        exgcd(b, a % b, y, x);
        y -= x * (a / b);
    }
}
ll invf(ll x, ll m){
    ll ans, y;
    exgcd(x, m, ans, y);
    ans = (ans % m + m) % m;
    return ans;
}
vector<int> p;
int tmp[maxn];
void mul(int *arr, int *brr, int *ans){
    for(auto it: p) tmp[it] = arr[brr[it]];
    for(auto it: p) ans[it] = tmp[it]; 
}
void qpow(ll ti){  
    for(; ti; ti >>= 1){
        if(ti & 1) mul(ans, arr, ans);
        mul(arr, arr, arr);
    }
}
void dfs(int now, int cnt){
    if(vis[now]){
        ll ti = invf(k, cnt);
        qpow(ti);
        p.clear();
    }else{
        p.push_back(now);
        vis[now] = 1;
        dfs(arr[now], cnt + 1);
    }
}
int main(void){
#ifdef ljxtt
freopen("data.in", "r", stdin);
#endif
    scanf("%d%d", &n, &k);
    for(int i = 1; i <= n; i++) scanf("%d", &arr[i]);
    for(int i = 1; i <= n; i++) ans[i] = i;
    for(int i = 1; i <= n; i++) if(!vis[i]) dfs(i, 0);
    for(int i = 1; i <= n; i++) printf("%d%c", ans[i], " \n"[i == n]);
    return 0;
}