Numpy 多维子集

我花了很长时间努力寻找这个问题的答案，但没有找到任何完全符合要求的答案。我有一个多维 numpy 数组，其中包含数据（在我的例子中是 3 维）和另一个数组（2 维），其中包含有关我想要沿原始数组的最后一个维度的值的信息。例如，这是一个说明问题的简单示例。我有一个数组a数据和另一个数组b包含沿维度 2 的索引a。我想要一个新的二维数组c where 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索引功能？

当您使用多维数组对多维数组进行花式索引时，每个维度的索引将一起广播。考虑到这一点，您可以这样做：

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