如何在NumPy中高效地对多轴应用排列数组？

给定排列数组perm：

import numpy as np
perm = np.array([[0, 1, 2], [1, 2, 0], [0, 2, 1], [2, 1, 0]])

单轴应用排列可以直接用v[perm]实现，但双轴应用时，原有的循环+meshgrid方法效率低下，这里提供一种完全向量化的高效实现方式：

实现代码 #

首先定义目标数组v：

v = np.arange(9).reshape(3, 3)

直接通过维度扩展+广播实现双轴索引：

# 向量化双轴排列实现
result = v[perm[:, :, np.newaxis], perm[:, np.newaxis, :]]

验证正确性 #

对比原方法的输出：

# 原低效实现
original_result = np.array([v[tuple(np.meshgrid(p, p, indexing="ij"))] for p in perm])

# 验证结果一致
print(np.array_equal(result, original_result))  # 输出：True

原理说明 #

• 通过np.newaxis为perm扩展维度，让行索引（perm[:, :, None]，形状(4,3,1)）和列索引（perm[:, None, :]，形状(4,1,3)）通过广播自动匹配成(4,3,3)的索引矩阵，正好对应每个排列对应的双轴索引。
• 完全依赖NumPy的底层向量化运算，避免了Python循环的开销，在处理大规模数组和大量排列时，速度和内存效率都会显著提升。

内容的提问来源于stack exchange，提问作者sams-studio