使用列索引数组索引 2D numpy 数组答案

【问题标题】：Indexing a 2D numpy array with an array of column indices使用列索引数组索引 2D numpy 数组
【发布时间】：2023-03-30 13:40:01
【问题描述】：

假设我有一个形状为 (N, D) 的二维 numpy 数组 A 和一个形状为 (N,) 的一维数组 b。我想使用第二个索引第一个数组以生成以下数组：

arr[0] = A[0][b[0]]
arr[1] = A[1][b[1]]
arr[2] = A[2][b[2]]
...
arr[N-1] = A[N-1][b[N-1]]

列表理解很容易做到这一点，但我希望有一个有效的矢量化解决方案，使用 numpy 方法或索引。

【问题讨论】：

【解决方案1】：

你可以使用numpy的integer-array-indexing

arr = A[range(N), b]

【讨论】：

【解决方案2】：

试试这个：

for i,x in enumerate(a): #loop iterate from a[0] to a[N-1]
  arr[i] = x[b[i]] #b[i] is from b[0] to b[N-1]

【讨论】：