NumPy：如何根据索引矩阵对二维矩阵进行行级排序？

一维数组的基础索引用法 #

在NumPy中，对一维数组按自定义索引重排是非常直接的操作：

>>> import numpy as np

>>> a = np.array([0, 100, 200])
>>> order = np.array([1, 2, 0])
>>> print(a[order])
[100 200   0]

二维索引的常见误区 #

当原数组和索引矩阵都是二维结构时，直接用数组索引的方式无法得到预期效果：

>>> A = np.array([[0, 1, 2],
...               [3, 4, 5],
...               [6, 7, 8]])

>>> Ord = np.array([[1, 0, 2],
...                 [0, 2, 1],
...                 [2, 1, 0]])

>>> print(A[Ord].shape)
(3, 3, 3)

出现这个结果的原因是NumPy整数索引的匹配规则：传入的二维索引矩阵Ord会被默认匹配到原数组A的第一维（行维度），Ord里每个元素的值都会被当成A的行号，取出对应的一整行数据。相当于每个索引位置都映射了A的整行内容，最终自然形成了三维的输出结果，和我们的预期不符。

我们真正想要的排序效果，是对A的每一行，单独应用Ord中对应的行索引做列位置的重排，最终输出和A形状完全一致的结果：

array([[1, 0, 2],
       [3, 5, 4],
       [8, 7, 6]])

正确实现方法 #

要实现逐行按列索引重排，需要生成对应的行索引，和列索引Ord做配对取值：

# 生成和Ord同形状的逐行行索引：每一行固定对应A的原行号
row_index = np.arange(A.shape[0])[:, None]
# 行索引与列索引广播配对，完成精准取值
result = A[row_index, Ord]

>>> print(result)
[[1 0 2]
 [3 5 4]
 [8 7 6]]

实现原理 #

• 生成的row_index是形状为(3,1)的二维数组，内容为[[0], [1], [2]]
• 这个行索引会和形状为(3,3)的列索引Ord自动广播，配对成一一对应的(行号,列号)坐标对，精准取出每个位置需要的元素，最终输出和原数组形状完全一致的重排结果。

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