E - Permute K times

Problem Statement

You are given a sequence $X$ of length $N$ where each element is between $1$ and $N$, inclusive, and a sequence $A$ of length $N$.
Print the result of performing the following operation $K$ times on $A$.

• Replace $A$ with $B$ such that $B_i = A_{X_i}$.

Constraints

• All input values are integers.
• $1 \le N \le 2 \times 10^5$
• $0 \le K \le 10^{18}$
• $1 \le X_i \le N$
• $1 \le A_i \le 2 \times 10^5$

Input

The input is given from Standard Input in the following format:

$N$ $K$
$X_1$ $X_2$ $\dots$ $X_N$
$A_1$ $A_2$ $\dots$ $A_N$

Output

Let $A'$ be the sequence $A$ after the operations. Print it in the following format:

$A'_1$ $A'_2$ $\dots$ $A'_N$

Sample Input 1

7 3
5 2 6 3 1 4 6
1 2 3 5 7 9 11

Sample Output 1

7 2 3 5 1 9 3

In this input, $X=(5,2,6,3,1,4,6)$ and the initial sequence is $A=(1,2,3,5,7,9,11)$.

• After one operation, the sequence is $(7,2,9,3,1,5,9)$.
• After two operations, the sequence is $(1,2,5,9,7,3,5)$.
• After three operations, the sequence is $(7,2,3,5,1,9,3)$.

Sample Input 2

4 0
3 4 1 2
4 3 2 1

Sample Output 2

4 3 2 1

There may be cases where no operations are performed.

Sample Input 3

9 1000000000000000000
3 7 8 5 9 3 7 4 2
9 9 8 2 4 4 3 5 3

Sample Output 3

3 3 3 3 3 3 3 3 3

解题思路

  比赛时一直想着线性的做法，完全没考虑过倍增，结果发现代码巨难写。

  首先如果按照 $i \to x_i$ 进行连边，就会得到一棵基环树，当 $k$ 足够大的时候其实就是在环上不断循环。对于每个点可以先让 $k$ 减去该点到环的距离，然后对环的大小取模，求出最后停留在环上的哪个位置，但实现起来非常复杂。本质上就是快速求出每个点在走 $k$ 步后最后停留在哪里，由于 $k$ 非常大这启发我们可以往倍增上考虑。

  定义 $X^{j}_{i}$ 表示从点 $i$ 经过 $2^j$ 步后到达的位置，$X^{0}_{i}$ 就是初始输入的数据。想要求得 $X^{j}_{i}$，我们可以先让 $i$ 走 $2^{j-1}$ 步到达 $X^{j-1}_{i}$，再从 $X^{j-1}_{i}$ 走 $2^{j-1}$ 步到达 $X^{j}_{i}$，因此有转移方程 $X^{j}_{i} = X^{j-1}_{X^{j-1}_{i}}$。

  为了求出从 $i$ 走 $k$ 步的结果，先把 $k$ 看作二进制数 $k_{59}k_{58} \cdots k_{0}, \, k_j \in \{0,1\}$，维护一个数组 $p_i$ 表示从 $i$ 经过若干步后到达的位置，初始时定义 $p_i = i$。然后枚举 $k$ 的每一个二进制位，如果 $k_j = 1$，意味着我们需要走 $2^{j}$ 步，更新 $p_i = X^{j}_{p_i}$。

  最后输出 $a_{p_i}$ 即可。

  AC 代码如下，时间复杂度为 $O(n \log{k})$：

#include <bits/stdc++.h>
using namespace std;

typedef long long LL;

const int N = 2e5 + 5;

int x[60][N], p[N], a[N];

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int n;
    LL m;
    cin >> n >> m;
    for (int i = 1; i <= n; i++) {
        cin >> x[0][i];
    }
    for (int i = 1; 1ll << i <= m; i++) {
        for (int j = 1; j <= n; j++) {
            x[i][j] = x[i - 1][x[i - 1][j]];
        }
    }
    iota(p + 1, p + n + 1, 1);
    for (int i = 0; 1ll << i <= m; i++) {
        if (m >> i & 1) {
            for (int j = 1; j <= n; j++) {
                p[j] = x[i][p[j]];
            }
        }
    }
    for (int i = 1; i <= n; i++) {
        cin >> a[i];
    }
    for (int i = 1; i <= n; i++) {
        cout << a[p[i]] << ' ';
    }
    
    return 0;
}

参考资料

  Editorial - AtCoder Beginner Contest 367：https://atcoder.jp/contests/abc367/editorial/10707