如何用Numpy labels索引数组提取二维data数组每行对应位置元素

实现方案 #

最简洁高效的矢量化实现可以直接用Numpy的整数数组高级索引特性，一行代码即可完成：

result = data[np.arange(len(labels)), labels]

原理解释 #

Numpy的整数数组索引规则为：当传入两个形状一致的整数数组分别作为行、列维度的索引时，会逐位置配对坐标提取元素：

• 行索引为np.arange(len(labels))，也就是[0, 1]，正好对应labels的下标也就是data的行号
• 列索引就是labels本身[3, 1]
• 最终会配对提取(0,3)和(1,1)两个位置的元素，正好得到目标结果[0.2, 0.3]

方案优势 #

• 完全矢量化实现，无Python层级循环，大数组下性能极高
• 语法符合Numpy官方设计规范，可读性强易维护
• 适配任意长度的labels输入，只要labels元素取值不超过data的最大列索引即可

完整可运行示例 #

import numpy as np

data = np.array([[0.75, 0.05, 0.1, 0.2],
                 [0.4, 0.3, 0.2, 0.1]])
labels = np.array([3, 1])

result = data[np.arange(len(labels)), labels]
print(result) # 输出：[0.2 0.3]

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