矩阵索引，为每一行提取不同的列

Question

底部有小数组的示例。 我有一个尺寸为 MxN (N>100) 的 numpy 数组数据和一个尺寸为 Mx100 的较小数组 A。

数组 A 包含数据数组的列索引。

我想提取由 A 中的索引给出的所有数据列。 问题是 A 中的索引对于每一行都是不同的。

我试过了

 data[:,A]

但这会返回一个 MxMx100 数组。 我只想要一个 Mx100 阵列。 我可以在 for 循环中执行此操作，但这需要很长时间。 我怎样才能正确地做到这一点？

例子：

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

我想提取： 第一行：索引 0 和 1。第二行：索引 1 和 2。第三行：索引 0 和 2

想要的结果：

array([[1,2],
[5,6],
[7,9]])

Answer 1

使用broadcasting

>>> data[np.arange(A.shape[0])[:,None], A]

array([[1, 2],
       [5, 6],
       [7, 9]])