如何使用二维索引数组按指定顺序排列NumPy数组元素

当然有简便方法啦！用NumPy的高级索引就能轻松搞定这个需求，而且代码简洁高效，完全不用写循环。

思路解析 #

你的数组a是一个形状为(2, 3, 3)的三维数组，可以理解为包含2个“块”，每个块是3行3列的二维数组。而M数组的每一行，正好对应每个块里要重新排列的行索引。我们只需要让每个块对应上M里的索引，就能完成重排。

代码实现 #

先构造示例数组来验证：

import numpy as np

# 构造示例数组a
a = np.array([
    [[a1, a2, a3], [a4, a5, a6], [a7, a8, a9]],
    [[b1, b2, b3], [b4, b5, b6], [b7, b8, b9]]
])

# 索引数组M
M = np.array([[1, 0, 2], [1, 2, 0]])

# 执行重排
result = a[np.arange(a.shape[0])[:, None], M]

为什么这样可行？ #

• np.arange(a.shape[0])生成了第一个维度的索引[0, 1]，代表我们要处理的2个块。
• [:, None]把这个一维数组变成形状为(2, 1)的二维数组，这样就能和M（形状(2, 3)）进行广播，确保每个块都能匹配到M中对应的索引行。
• 最终的索引组合会让NumPy对每个块，按照M里的索引顺序提取行，正好得到你想要的结果。

验证结果 #

运行上面的代码后，result的输出就是：

array([
    [[a4, a5, a6], [a1, a2, a3], [a7, a8, a9]],
    [[b4, b5, b6], [b7, b8, b9], [b1, b2, b3]]
])

完全符合你的需求！

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