NumPy中二维数组索引另一二维数组的操作效果解析

先明确基础前提：

• 被索引数组a是形状为(21, 3)的二维数组，第0轴（行轴）长度为21，第1轴（列轴）长度为3
• 索引数组b是形状为(225, 300)的二维数组，所有元素值都是合法的a第0轴索引（取值范围0~20，无越界）

逐段拆解执行逻辑 #

1. 首先看b[:]：这是对数组b做全维度切片，返回和b数据、形状完全一致的视图，在这里写不写[:]效果完全等价，a[b, :]和当前写法的运行结果没有任何区别。
2. 索引匹配规则：NumPy会把逗号分隔的索引按顺序对应到被索引数组a的各个轴上：
   • 第一个位置传入的b[:]是整数数组，对应a的第0轴（行轴），触发**高级整数索引（花式索引）**规则：b中的每一个元素值，都代表要从a的第0轴上取对应序号的行。
   • 第二个位置传入的:是普通切片，对应a的第1轴（列轴），代表取对应行的全部列，不做列维度筛选。
3. 输出形状规则：高级整数索引的输出数组形状，为「传入的整数索引数组的形状」拼接「被索引数组剩余未被整数索引的维度形状」。这里索引数组b形状为(225, 300)，a剩余的第1轴长度为3，因此最终输出数组的形状为(225, 300, 3)。
4. 元素对应关系：对于输出数组out，任意位置的元素满足 out[i, j, :] = a[b[i, j], :]——相当于遍历b的每一个位置，把位置上存的行号替换成a中对应行的整行内容。

实际作用 #

这个写法最常见的使用场景是索引映射查表：比如把a的每一行存为一个id对应的特征向量（比如21个类别的RGB三通道颜色值、21个类别的特征嵌入），b存的是网格状的id编号（比如225*300分辨率的图像分割类别图、网格点位的分类结果），一行代码就能完成全网格的id到对应特征的映射，不需要写Python层循环，运行是C层面的矢量计算，效率比手写循环高几个数量级。

极简验证示例：如果a = np.array([[0,0,0], [255,255,255]])（形状(2,3)，存黑、白两种颜色），b = np.array([[0,1],[1,0]])（形状(2,2)，存2*2网格的类别），那a[b, :]的输出为

[[[  0,   0,   0],
  [255, 255, 255]],
 [[255, 255, 255],
  [  0,   0,   0]]]

形状为(2,2,3)，正好是把类别图映射成了2*2分辨率的黑白RGB图。

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