如何在numpy中找到每行中最大值的索引，行的连接？

Question

我不知道这是否简单，或者之前是否有人问过。 （我搜索了但没有找到正确的方法。我找到了numpy.argmax和numpy.amax但我无法正确使用它们。）

我有一个 numpy 数组（它是一个CxKxN 矩阵）如下（C=K=N=3）：

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

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

       [[5, 2, 1],
        [3, 3, 3],
        [4, 1, 2]]])

我想找到每行中最大元素的索引。一行是每个矩阵的三行 (C) 的串联。换句话说，i-th 行是第一个矩阵中的i-th 行、第二个矩阵中的i-th 行...的串联，直到i-th C-th 矩阵中的行。

比如第一行是

[1, 2, 3, 2, 1, 1, 5, 2, 1]

所以我想退货

[2, 0, 0] # the index of the maximum in the first line

和

[0, 1, 2] # the index of the maximum in the second line

和

[0, 2, 0] # the index of the maximum in the third line

或

[1, 2, 1] # the index of the maximum in the third line

或

[2, 2, 0] # the index of the maximum in the third line

现在，我正在尝试这个

np.argmax(a[:,0,:], axis=None) # for the first line

它返回6 和

np.argmax(a[:,1,:], axis=None) 

它返回2 和

np.argmax(a[:,2,:], axis=None) 

它返回0

但我可以将这些数字转换为 6 = (2,0,0) 等索引。

Answer 1

通过转置和重塑，我得到了你的“行”

In [367]: arr.transpose(1,0,2).reshape(3,9)
Out[367]: 
array([[1, 2, 3, 2, 1, 1, 5, 2, 1],
       [2, 1, 4, 1, 3, 1, 3, 3, 3],
       [4, 3, 3, 3, 4, 2, 4, 1, 2]])
In [368]: np.argmax(_, axis=1)
Out[368]: array([6, 2, 0])

这些最大值与您的相同。相同的索引，但在 (3,3) 数组中：

In [372]: np.unravel_index([6,2,0],(3,3))
Out[372]: (array([2, 0, 0]), array([0, 2, 0]))

将它们加入中间维度范围：

In [373]: tup = (_[0],np.arange(3),_[1])
In [374]: np.transpose(tup)
Out[374]: 
array([[2, 0, 0],
       [0, 1, 2],
       [0, 2, 0]])