python - 使用来自另一个数组的索引有效地对数组进行切片

Question

（如果这是一个重复的问题，我提前道歉，虽然我在 SO 上查看了许多类似的问题，但没有找到匹配的解决方案）

假设你有一个数组

A = np.array([
    [0, 1, 2],
    [3, 4, 5],
    [6, 7, 8]
])

和另一个数组

I = np.array([1, 1, 2])

对于 中的每一行，A我想获取它的第i个元素，其中i是 的第行元素I。

在这种情况下，我想要的输出是array([1, 4, 8]).

我最直观的尝试是：

A[:, I]

然后我认为所需的输出实际上是它的对角线，所以A[:, I].diagonal()可以做到这一点。

但是感觉这样做会浪费一些空间和时间，因为它需要一个中间的“大”矩阵，从中提取对角线。

执行此切片是否有更有效的方法？

标签： pythonnumpy

Answer 1

这可以解决问题：

res = A[np.arange(A.shape[0]), I]