是否有 numpy 方法通过索引列表过滤 arrays 数组

Question

假设我有 numpy 数组x其中x是

x = np.array(
  [
    [1, 2],
    [3, 4]
  ]
)

我还有一个索引列表i其中i = [0, 1] 。 我想为x中的每个n获取索引n[i]处的值的数组。

最佳 output 是

np.array([1, 4])

很明显，这可以通过 for 循环来完成，但我想知道是否有更简单的 *numpy* 方法来做到这一点。

Answer 1

您可以使用索引：

x[np.arange(x.shape[0]), i]

output： array([1, 4])