python - Numpy 多维子集

Question

我一直在努力寻找这个问题的答案，但没有找到任何符合要求的答案。我有一个包含数据的多维 numpy 数组(在我的例子中是 3 维)和另一个数组(2 维)，它包含有关我想要沿原始数组的最后一个维度的值的信息。例如，这是一个说明问题的简单示例。我有一个数据数组 a，另一个数组 b 包含沿 a 的第 2 维的索引。我想要一个新的二维数组 c，其中 c[i, j] = a[i, j, b[i, j]]。我唯一能做的想用一个循环来做到这一点，如下所述。然而，这看起来笨拙且缓慢。

In [3]: a = np.arange(8).reshape((2, 2, 2))
In [4]: a
Out[4]: 
array([[[0, 1],
        [2, 3]],

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

In [6]: b = np.array([[0, 1], [1, 1]])

In [8]: c = np.zeros_like(b)

In [9]: for i in xrange(2):
   ...:     for j in xrange(2):
   ...:         c[i, j] = a[i, j, b[i, j]]

In [10]: c
Out[10]: 
array([[0, 3],
       [5, 7]])

是否有更 pythonic 的方式来执行此操作，也许是我不知道的一些 numpy 索引功能？

Answer 1

当你用多维数组对多维数组进行花式索引时，每个维度的索引都会一起广播。考虑到这一点，您可以:

>>> rows = np.arange(a.shape[0])
>>> cols = np.arange(a.shape[1])
>>> a[rows[:, None], cols, b]
array([[0, 3],
       [5, 7]])