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

【问题标题】：Numpy: 2D array access with 2D array of indicesNumpy：使用二维索引数组访问二维数组
【发布时间】：2014-02-10 13:40:25
【问题描述】：

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

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，我想避免 for 循环。

【问题讨论】：

您似乎在a 周围多了一个括号。 a = array([[0,0],[1,1],[2,0],[2,1]], dtype=int) 有效吗？
@JLLagrange 它应该在那里。 a.shape 应该是 (2,2,2)，或更一般地说，(n,m,2) 和结果 (n,m,1) (=(n,m))。

标签： python arrays numpy

【解决方案1】：

实际上，这是可行的：

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

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

【讨论】：

【解决方案2】：

对于这种情况，这是可行的

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

【讨论】：

需要重新调整结果：b.reshape(a.shape[:-1])

【解决方案3】：

更通用的解决方案，当您想要使用形状为 (n,m) 任意大尺寸 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]))

【讨论】：