NumPy如何无Python for循环向量化生成大量数组随机排列

原有实现的性能瓶颈来自Python层面的for循环，每次循环调用np.random.permutation都会产生额外的解释器开销，这类开销可以通过numpy原生的批量操作完全规避。

最优实现（NumPy 1.17+ 新版随机API） #

新版随机数生成器直接提供了沿指定轴批量打乱数组的接口，不需要手动实现排序逻辑，性能最高：

import numpy as np

arr = np.array([1, 1, 2, 2, 3, 3, 4, 4])
n_rows = 100_000

# 初始化默认随机数生成器
rng = np.random.default_rng()
# 先将原数组广播为目标形状的二维数组（无额外内存拷贝），再沿行方向批量打乱
result = rng.permuted(np.broadcast_to(arr, (n_rows, arr.size)), axis=1)

10万行规模下该实现的计时结果约为19ms，和np.random.randint生成同形状随机数组的速度几乎一致。

兼容旧版本NumPy的通用实现 #

如果使用低于1.17版本的numpy，可以通过随机权重排序的方式实现向量化，核心逻辑是为每个位置生成独立的随机值，按行排序得到乱序索引，再通过索引取原数组内容：

n_rows = 100_000
arr_len = arr.size
# 生成[0,1)区间的随机数矩阵，形状和目标输出一致
rand_vals = np.random.rand(n_rows, arr_len)
# 按行取argsort得到每行的乱序索引
shuffle_indices = rand_vals.argsort(axis=1)
# 广播索引原数组得到结果
result = arr[shuffle_indices]

该实现10万行规模下的计时结果约为23ms，相比原有1.06s的循环实现性能提升超过45倍，且性能随行数增长的下降幅度完全来自numpy底层C实现的线性开销，不会出现Python循环带来的额外性能陡降。

实现说明 #

• 两种实现生成的排列都满足均匀随机的要求，和单步调用np.random.permutation生成排列的概率分布一致。
• np.broadcast_to不会复制原数组数据，内存占用远低于先把原数组复制扩展为n行再打乱的实现。
• 不要使用np.random.shuffle做批量操作，该函数是原地操作且不支持指定轴参数，无法实现行维度的独立打乱。

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