如何在 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)

在这种情况下，输出仍然应该是一个(2, 2) 矩阵，但每个单元格都包含一个维数为 0 和维数为 1 的最大索引的元组，即

示例输出

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

---

解决方案

正如@hpaulj 正确指出的那样，这个问题的解决方案在于调用

idx = B.reshape(6,2,2).argmax(axis=0)
2d_indices = np.unravel_index(idx,(2,3))

但是，这种方式的排序与我的预期有点不同（并在我上面的问题中概述）。要得到与上面完全相同的输出，只需添加

list_of_tuples = [a for a in zip(2d_indices[0].flatten(), 2d_indices[1].flatten())]
np.array(list_of_tuples).reshape(3, 3, -1)

对我来说看起来有点麻烦，可能有更直接的方法可以到达那里，但为了完整起见，我仍然想发布它:)

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