Python 使用来自另一个数组的索引高效地切片一个数组

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

假设您有一个数组

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

和另一个数组

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

对于

A

中的每一行，我想得到它的第I个元素，其中I是

I

的第th行元素

在本例中，我想要的输出是

数组（[1,4,8]）

我最直观的尝试是：

A[：，I]

然后我计算出所需的输出实际上是它的对角线，因此

A[：，I].diagonal（）

就可以了

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

---

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

这将实现以下目的：

res=A[np.arange（A.shape[0]），I]

这正是我想要的。与我之前使用的.diagonal（）方法相比，它的加速和内存节省是巨大的。非常感谢你。