python - 如何获取 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 仅接受标量轴值(其他一些函数允许元组)。但我们可以 reshape 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]]])