在Numpy数组的第3维上使用Numpy armmax[重复]

Question

我有一个3d numpy数组的形状（150，9，5），我想得到第三维的argmax（从元素中），但我无法完成。例如，假设数组如下所示（此示例不遵循我在代码中使用的形状）：

[[[615, 142, 18]],

   [[12, 412, 98]],

   [[5, 11, 162]],

   [[19, 76, 7]],

   [[56, 78, 12]],

   [[119, 41, 55]],

   [[26, 7, 15]],

   [[19, 67, 53]]]

我想得到一个形状数组（150,9,1）（这也是我的情况，与示例无关）。例如，它将是：

[[[0]],

   [[1]],

   [[2]],

   [[1]],

   [[1]],

   [[0]],

   [[0]],

   [[1]]]

当我用0轴和1轴尝试np.argmax（）时，我得到了错误的结果。

有没有办法直接计算出来，或者我应该使用for循环遍历每个（9,5）？

Answer 1

您有额外的轴和额外的尺寸：

import numpy as np
a=np.array([[[615, 142, 18]],

   [[12, 412, 98]],

   [[5, 11, 162]],

   [[19, 76, 7]],

   [[56, 78, 12]],

   [[119, 41, 55]],

   [[26, 7, 15]],

   [[19, 67, 53]]])

a.argmax(axis=2)

输出：

Out[82]: 
array([[0],
       [1],
       [2],
       [1],
       [1],
       [0],
       [0],
       [1]], dtype=int64)