python - Numpy:使用 2D 索引数组访问 2D 数组

Question

我有两个数组，一个是索引对矩阵，

a = array([[[0,0],[1,1]],[[2,0],[2,1]]], dtype=int)

另一个是在这些索引处访问的数据矩阵

b = array([[1,2,3],[4,5,6],[7,8,9]])

并且我希望能够使用 a 的索引来获取 b 的条目。只是在做:

>>> b[a]

不起作用，因为它为 a 中的每个条目提供一行 b，即

array([[[[1,2,3],
         [1,2,3]],

        [[4,5,6],
         [4,5,6]]],


       [[[7,8,9],
         [1,2,3]],

        [[7,8,9],
         [4,5,6]]]])

当我想使用a最后一个轴中的索引对来给出b的两个索引时:

array([[1,5],[7,8]])

是否有一种干净的方法可以做到这一点，或者我是否需要 reshape b 并以相应的方式组合 a 的列？

在我的实际问题中，a 有大约 500 万个条目，b 是 100×100，我想避免 for 循环。

Answer 1

实际上，这有效:

b[a[:, :, 0],a[:, :, 1]]

给出 array([[1, 5], [7, 8]]).