如何用Numpy实现按行指定索引提取数组元素？

给定以下Numpy数组：

import numpy as np
a = np.array([[1,2,3],[4,5,6]])
idx = np.array([[0,2],[1,2]])

需要实现第i行用idx第i行作为列索引提取元素的效果，等价于循环实现：

np.vstack([a[i, idx[i]] for i in range(len(a))])

期望输出：

[[1,3],
 [5,6]]

原生Numpy解决方案 #

可以通过构造匹配形状的行索引数组，结合idx实现：

result = a[np.arange(a.shape[0])[:, None], idx]

原理说明 #

• np.arange(a.shape[0])生成与a行数对应的行索引数组：array([0, 1])
• [:, None]将行索引数组转换为列向量（形状变为(2,1)），与idx的(2,2)形状匹配后触发Numpy广播机制
• 此时行索引与列索引逐行对应：第0行取idx[0]指定的列，第1行取idx[1]指定的列，最终得到目标二维数组

错误方法分析 #

• a[idx]：直接将idx的值作为行索引，提取a的第0、2行和第1、2行，完全不符合需求
• a[:, idx]：对a的每一行都取idx中所有列索引，得到形状为(2,2,2)的三维数组，并非目标结果

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