如何在 4d numpy 數組中沿 2d 獲取最大值的索引

Question

想象一下，您有一個 3d 陣列，例如

A = np.array([[[1,2], [2,3] ], [[3,4], [1,2]], [[2,3], [3,4]]])

A.shape (3, 2, 2)

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

       [[3, 4],
        [1, 2]],

       [[2, 3],
        [3, 4]]])

現在，如果我想獲得沿第一維的最大值的索引，這很容易，因為

A.argmax(axis=0)

array([[1, 1],
       [2, 2]])

我如何對 4d 數組執行相同的操作，找到沿前兩個維度的最大值？

示例 4d 數組

np.random.randint(0, 9,[2,3,2,2])

B = array([[[[5, 4],
         [3, 8]],

        [[3, 5],
         [0, 4]],

        [[0, 1],
         [3, 0]]],


       [[[0, 2],
         [7, 3]],

        [[7, 3],
         [8, 0]],

        [[8, 3],
         [2, 7]]]])

B.shape (2, 3, 2, 2)

在這種情況下，output 應該仍然是一個(2, 2)矩陣，但是每個單元格都包含一個維度為 0 和維度為 1 的最大索引的元組，即

示例 output

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

Answer 1

argmax只接受標量軸值（其他一些函數允許元組）。 但我們可以重塑B ：

In [18]: B.reshape(6,2,2).argmax(axis=0)
Out[18]: 
array([[5, 1],
       [4, 0]])

並將它們轉換回二維索引：

In [21]: np.unravel_index(_18,(2,3))
Out[21]: 
(array([[1, 0],
        [1, 0]]),
 array([[2, 1],
        [1, 0]]))

這些值可以以各種方式重新排序，例如：

In [23]: np.transpose(_21)
Out[23]: 
array([[[1, 2],
        [1, 1]],

       [[0, 1],
        [0, 0]]])