Numpy：使用二维索引数组访问二维数组

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]

不工作，因为它给出了B的一行中的每个条目a ，即

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]])

是否有这样做的一个干净的方式，或者我需要重塑b和组合的列a以相应的方式？

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

Answer 1

实际上，这是有效的：

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

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

Answer 2

对于这种情况，这有效

tmp =  a.reshape(-1,2)
b[tmp[:,0], tmp[:,1]] 

Answer 3

一个更通用的解决方案，每当您想要使用具有任意大维度 m的形状 (n,m) 索引的二维数组，命名为inds ，以便访问另一个二维形状数组 (n,k) 的元素，命名为B ：

# array of index offsets to be added to each row of inds
offset = np.arange(0, inds.size, inds.shape[1])

# numpy.take(B, C) "flattens" arrays B and C and selects elements from B based on indices in C
Result = np.take(B, offset[:,np.newaxis]+inds)

另一个不使用np.take并且我觉得更直观的解决方案如下：

B[np.expand_dims(np.arange(B.shape[0]), -1), inds]

这种语法的优点是它既可以用于基于inds （如np.take ）从B读取元素，也可以用于赋值。

您可以使用以下方法进行测试，例如：

B = 1/(np.arange(n*m).reshape(n,-1) + 1)
inds = np.random.randint(0,B.shape[1],(B.shape[0],B.shape[1]))