如何在NumPy中根据指定向量按行选取对应位置元素？

假设存在矩阵A：

A = [[34, 61, 29, 74, 17, 32, 72, 92, 93, 57],
     [46, 10, 23, 84, 74, 57, 56, 88, 90, 36],
     [23, 83, 58, 42, 93, 54, 82, 48, 63, 73]]

以及长度为3的向量b：

b = [4, 0, 1]

是否存在NumPy函数可以实现如下功能：

>>> A.choose_row_wise(b)
输出结果：
[17, 46, 83]

有，NumPy无需专门的choose_row_wise方法，通过原生索引或内置函数就能实现需求，以下是两种高效实现方式：

• 基础索引法：这是NumPy中最常用的实现方式，性能最优

  import numpy as np
  
  A = np.array([[34, 61, 29, 74, 17, 32, 72, 92, 93, 57],
                [46, 10, 23, 84, 74, 57, 56, 88, 90, 36],
                [23, 83, 58, 42, 93, 54, 82, 48, 63, 73]])
  b = np.array([4, 0, 1])
  
  result = A[np.arange(A.shape[0]), b]
  print(result)  # 输出: [17 46 83]
  

  其中np.arange(A.shape[0])生成与行数匹配的行索引[0,1,2]，和b的列索引一一对应，精准定位每行要提取的元素。

• take_along_axis函数法：语义更清晰，明确指定按行方向选择对应位置的元素

  result = np.take_along_axis(A, b[:, np.newaxis], axis=1).flatten()
  print(result)  # 输出: [17 46 83]
  

  b[:, np.newaxis]将一维的b转为列向量以匹配A的维度，take_along_axis按指定轴（axis=1即列方向）提取对应元素，最后用flatten()转回一维数组。

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