如何在n维数组中选择值

Question

我一直在尝试执行一个简单的操作，但是我似乎找不到一种简单的方法来使用Numpy函数而不创建不必要的数组副本。

假设我们具有以下3维数组：

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

我目前唯一的解决方案是在x和y的第一维上使用for循环，如下所示：

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

有没有办法避免使用numpy函数或奇特的索引进行for循环？

谢谢！

Answer 1

棘手的方面是，您不希望每个切片都具有所有的第0维，而是希望这些切片与第0维中的每个元素相对应 。 因此，您可以执行以下操作：

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

Answer 2

花式索引：

x[np.arange(y.size),:,y]

得到：

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