使用較小的N-D陣列按列對N-D numpy數組進行排序

Question

據我所知，通過使用花式索引對Sort ND numpy數組進行使用花式索引，我可以執行以下c = a[:, :, b]其中b定義了我想按列排序的順序

>>> a = np.array([[[ 0,  1], [ 2,  3]],
                  [[ 4,  5], [ 6,  7]],
                  [[ 8,  9], [10, 11]]])
>>> b = np.array([1, 0])
>>> c = a[:, :, b]
>>> c
array([[[ 1,  0],
        [ 3,  2]],

       [[ 5,  4],
        [ 7,  6]],

       [[ 9,  8],
        [11, 10]]])

現在我用b2兩個輸入增加b ，對應於我想要在a對每組2x2進行排序a

>>> b2 = np.array([[1, 0], [0, 1], [1, 0]])
>>> c2 = ?
>>> c2
array([[[ 1,  0],
        [ 3,  2]],

       [[ 4,  5],
        [ 6,  7]],

       [[ 9,  8],
        [11, 10]]])

我有一組更大的輸入，我有一個函數返回一個類似'b2'的數組，它為我提供了我應該獲得的信息。 因此，我可以知道我應該填寫什么c2 = ? 為了得到理想的結果？

Answer 1

這是一種花式索引的方法 -

(a[np.arange(a.shape[0])[:,None],:,b2]).transpose(0,2,1)

樣品運行 -

In [191]: a
Out[191]: 
array([[[7, 8, 5, 2, 0],
        [6, 7, 0, 7, 1],
        [7, 6, 5, 4, 0]],

       [[8, 0, 5, 5, 7],
        [4, 3, 4, 0, 1],
        [8, 6, 3, 2, 4]],

       [[3, 2, 7, 3, 7],
        [4, 3, 0, 1, 5],
        [4, 3, 7, 8, 7]]])

In [192]: b2
Out[192]: 
array([[1, 2, 4, 3, 0],
       [4, 2, 0, 1, 3],
       [1, 3, 4, 0, 2]])

In [193]: (a[np.arange(a.shape[0])[:,None],:,b2]).transpose(0,2,1)
Out[193]: 
array([[[8, 5, 0, 2, 7],
        [7, 0, 1, 7, 6],
        [6, 5, 0, 4, 7]],

       [[7, 5, 8, 0, 5],
        [1, 4, 4, 3, 0],
        [4, 3, 8, 6, 2]],

       [[2, 3, 7, 3, 7],
        [3, 1, 5, 4, 0],
        [3, 8, 7, 4, 7]]])

Answer 2

如果沒有人找到純粹的花式索引解決方案，這里就是循環第一個軸的解決方案：

np.asarray([a[n,:,p] for n,p in enumerate(b2)])

array([[[ 1,  3],
    [ 0,  2]],

   [[ 4,  6],
    [ 5,  7]],

   [[ 9, 11],
    [ 8, 10]]])

Answer 3

類似於@Divakar's解決方案，但沒有transpose 。

In [259]: I,J, K = np.ogrid[:3,:2,:2]

In [260]: a[I, J, b[:,None,:]]
Out[260]: 
array([[[ 1,  0],
        [ 3,  2]],

       [[ 4,  5],
        [ 6,  7]],

       [[ 9,  8],
        [11, 10]]])

我使用ogrid （或np.ix_ ）作為生成2個3d數組的緊湊方式，用b[:,None,:]廣播以產生一組(3,2,2)索引。

完全Nones的等價物是：

a[np.arange(3)[:,None,None], np.arange(2)[None,:,None], b[:,None,:]]

（這更清楚地表明b正沿着第一軸和最后一軸選擇項目）

要查看完整的廣播索引數組，請打印：

 np.broadcast_arrays(I,J,b[:,None,:])