python - 使用二维数组作为索引从二维 numpy 数组中提取元素

Question

这个问题已经有答案了:

Indexing one array by another in numpy (4 个回答)

已关闭 5 年前。

我有一个二维数组:

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

我还有这个数组，它的列数与输入数组相同，并且每行都包含从输入数组的匹配列中提取的索引。

>>> indices = np.array([[0],[1]])
>>> indices
array([[0],
       [1]])

在此示例中，我希望获得以下数组作为输出:

array([[1],
       [4]])

有什么办法可以实现这个目标吗？

Answer 1

通过 NumPy 数组索引的一种方法。

A = np.array([[1,2],[3,4]])
idx = np.array([[0],[1]])

res = A[np.arange(A.shape[0])[:, None], idx]

print(res)

[[1]
 [4]]

请注意，我们需要确保行索引器和列索引器具有相同的形状，在本例中为 (n, 1)，其中 n 是行数。