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

有没有干净的方法，或者我是否需要重新塑造b并以相应的方式合并a的列？

在我的实际问题a中有大约500万个条目，而b是100分100，我希望避免循环。

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

更通用的解决方案，只要您想使用形状（n，m）的二维数组，任意大维度m ，名为 inds ，按顺序访问另一个形状为（n，k）的2D数组的元素，命名为 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)

您可以使用例如：

进行测试

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

Numpy：具有2D索引数组的2D数组访问

3 个答案: