Numpy二维数组用二维数组索引报错及同类场景差异解析

第一个报错的原因 #

A是形状为(2, 4)的二维数组，当你直接用A[B]时，numpy的索引规则是：用多维索引数组索引多维数组时，默认会把索引数组的每个元素对应到原数组的最外层轴（axis=0）。

B里的元素是[2,1,0,3]和[2,0,3,1]，这些数值被用来索引A的axis=0（也就是行），但A的axis=0只有0、1两个合法索引，2和3显然超出范围，所以抛出IndexError。

你想要的是对每行的列进行排序，所以需要明确指定行索引和列索引的对应关系，不能只给列索引。

第二个案例可行的原因 #

C是形状为(4,)的一维数组，numpy对一维数组的索引规则更直接：不管索引数组是几维，每个元素都会被当作一维数组的轴索引来用。B里的所有元素（0-3）都在C的合法索引范围内，所以C[B]会生成一个(2,4)的数组，每个位置对应C中该索引的元素，完全合法。

正确实现A按argsort结果排序的写法 #

要实现对A的每行按B的索引排序，需要给每行匹配对应的行索引，利用numpy的广播机制：

import numpy as np
A = np.array([
    [10, 7, 3, 100],
    [1, 7, 0, 3],
])
B = np.argsort(A, axis=1)
# 生成行索引：(2,1)形状，和B(2,4)广播匹配
row_indices = np.arange(A.shape[0])[:, None]
sorted_A = A[row_indices, B]
# sorted_A结果：
# [[  3   7  10 100]
#  [  0   1   3   7]]

这里row_indices是[[0], [1]]，和B广播后，每个行索引对应每行的列索引，numpy会正确定位到A[0,2]、A[0,1]等元素，实现每行按排序索引取数。

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