如何在n维数组中选择值

Question

我一直在尝试执行一个简单的操作,但我似乎无法找到一个简单的方法来使用Numpy函数来完成它而不创建不必要的数组副本.

假设我们有以下三维数组:

In [171]: x = np.arange(24).reshape((4, 3, 2))
In [172]: x
Out[172]: 
array([[[ 0,  1],
        [ 2,  3],
        [ 4,  5]],

       [[ 6,  7],
        [ 8,  9],
        [10, 11]],

       [[12, 13],
        [14, 15],
        [16, 17]],

       [[18, 19],
        [20, 21],
        [22, 23]]])

以下数组:

In [173]: y = np.array([0, 1, 1, 0])

我想x为每一行选择索引是相应元素的最后一个维度的值y.换句话说,我想:

array([[ 0,  2, 4],
       [ 7,  9, 11],
       [13, 15, 17],
       [18, 20, 22]])

我目前拥有的唯一的解决办法是使用一个for循环在第一维x和y,如下所示:

z = np.zeros((4, 3), dtype=int)
for i, row in enumerate(x):
    z[i, :] = row[:, y[i]]

有没有办法在这里避免使用for循环,使用numpy函数或花哨的索引？

谢谢!

Answer 1

棘手的方面是您不希望每个切片都具有所有第0维,您希望切片对应于第0维中的每个元素.所以你可以这样做:

>>> x[np.arange(x.shape[0]), :, y]
array([[ 0,  2,  4],
       [ 7,  9, 11],
       [13, 15, 17],
       [18, 20, 22]])