使用列表選擇多維數組的一維 - 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

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必須以一種或其他方式映射或廣播到其他維度。

Answer 2

您還需要“逐行”選擇其余維度。 也就是說，您需要提供以下形式的索引數組： 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])