用列表选择一维多维数组 - numpy

Question

我有一个形状为 (800,5,4) 的 3D 数组，例如：

arr = array([[35.  , 33.  , 33.  ,  0.15],
       [47.  , 47.  , 44.  ,  0.19],
       [49.  , 56.  , 60.  ,  0.31],
       ...,
       [30.  , 27.  , 25.  ,  0.07],
       [54.  , 49.  , 42.  ,  0.14],
       [33.  , 30.  , 28.  ,  0.22]])

我有一个第二维的一维数组（因此它们的范围从 0 到 4），如下所示：

indeces = [0,3,2,0,1,1,1,0,...,0,1,2,2,4,3]

我想从第二维中选择idx项，并取回一个形状为(800,4)的数组

我尝试了以下方法，但无法成功：

indexed = arr[:,indeces,:]

我错过了什么？

Answer 1

您还需要“逐行”选择剩余的维度。也就是说，您需要提供以下形式的索引数组：range(arr.shape[0])。例如：

>>> a = np.arange(27).reshape(3, 3, 3)
>>> i = [0, 1, 2]
>>> a[range(a.shape[0]), i, :]
array([[ 0,  1,  2],
       [12, 13, 14],
       [24, 25, 26]])
>>> a[range(a.shape[0]), i, range(a.shape[2])]
array([ 0, 13, 26])

Answer 2

In [178]: arr = np.arange(24).reshape(2,3,4)    

如果我有一个包含 7 个项目的列表：

In [179]: idx = [0,1,1,2,2,0,1]                                                                  
In [180]: arr[:,idx,:]                                                                           
Out[180]: 
array([[[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11],
        [ 8,  9, 10, 11],
        [ 0,  1,  2,  3],
        [ 4,  5,  6,  7]],

       [[12, 13, 14, 15],
        [16, 17, 18, 19],
        [16, 17, 18, 19],
        [20, 21, 22, 23],
        [20, 21, 22, 23],
        [12, 13, 14, 15],
        [16, 17, 18, 19]]])
In [181]: _.shape                                                                                
Out[181]: (2, 7, 4)

为了产生 (2,4) 的结果，我们必须从其他维度为每一对在第二个暗淡上选择一个元素。

一般情况是使 idx 成为 (2,4) 数组，并使用也广播到 (2,4) 的维度进行索引：

In [182]: idx = np.array([0,1,1,2,2,0,1,0]).reshape(2,4)                                         
In [183]: arr[np.arange(2)[:,None],idx,np.arange(4)]                                             
Out[183]: 
array([[ 0,  5,  6, 11],
       [20, 13, 18, 15]])
In [184]: _.shape                                                                                
Out[184]: (2, 4)

或者我们可以选择一个标量：

In [185]: arr[:,2,:]                                                                             
Out[185]: 
array([[ 8,  9, 10, 11],
       [20, 21, 22, 23]])

@a_guest 展示了如何使用与第一个维度匹配的idx 来执行此操作（并切片最后一个维度）。

您的idx 必须以其他方式映射或广播其他维度。