如何高效对NumPy数组多次应用置换索引操作？

我有一个尺寸为N的数组A，以及一个对应数组A的置换索引数组B。我希望对数组A多次应用该置换操作，请问是否有简洁高效的NumPy命令可以实现？

可复现示例代码：

import numpy as np
np.random.seed(121421)

# Example Array
A = np.random.uniform(0, 100, (100, 3))

# Permutation indices
B = np.random.choice(np.arange(A.shape[0]), A.shape[0], replace=False)

# Permute many times
npermut = 10000
C = np.array(B)
for n in range(npermut-1):
    C = C[B]
print(A[C])

直接循环执行k次置换的时间复杂度为O(k*N)，当k数值较大时（比如示例中的10000）效率极低。利用置换的循环分解特性可以将时间复杂度降至O(N)，无论k多大都能快速得到结果，具体实现如下：

核心原理 #

任何置换都可以分解为若干个独立的循环。例如，若元素i所在的循环长度为m，那么经过m次置换后i会回到初始位置。因此，经过k次置换后，i的位置等价于循环内右移k % m位后的位置。

实现代码 #

import numpy as np
np.random.seed(121421)

A = np.random.uniform(0, 100, (100, 3))
B = np.random.choice(np.arange(A.shape[0]), A.shape[0], replace=False)
npermut = 10000

# 1. 标记已访问的元素，避免重复处理
visited = np.zeros(len(B), dtype=bool)
final_indices = np.arange(len(B))

# 2. 遍历所有元素，分解置换循环
for i in range(len(B)):
    if not visited[i]:
        # 提取当前循环的所有元素
        cycle = []
        j = i
        while not visited[j]:
            visited[j] = True
            cycle.append(j)
            j = B[j]
        # 3. 计算当前循环经过npermut次置换后的位置映射
        cycle_len = len(cycle)
        shift = npermut % cycle_len
        # 循环右移shift位，得到置换后的元素位置
        shifted_cycle = cycle[-shift:] + cycle[:-shift]
        # 更新最终索引数组
        final_indices[cycle] = shifted_cycle

# 4. 直接应用最终索引得到结果
result = A[final_indices]

效率对比 #

• 原循环方法：需要执行npermut-1次索引查找，当npermut=10000时，需执行9999次数组索引操作。
• 循环分解方法：仅需一次遍历完成循环分解，后续直接生成最终索引，运行时间与npermut无关，大幅提升效率。

内容的提问来源于stack exchange，提问作者yvrob