获取每行numpy数组中的最高值的索引？

Question

我有一个30行4列的numpy数组，对于每一行，我需要获取最大值所在的索引。

所以对于这样的数组

a = np.array([[0, 1, 2],[7, 4, 5]])

我想获得一个列表，其中第一行的索引为2，第二行的索引为0。

我尝试使用numpy函数argmax，如下所示：

for i in range(len(a)):
    results=[np.argmax(a)]
return (results)

但是我只获得全局最大值，有人知道如何解决这个问题吗？

非常感谢您的帮助。

Answer 1

将argmax方法与axis=1配合使用以处理行。

>>> import numpy as np
>>> a = np.array([[0, 1, 2],[7, 4, 5]])
>>> a.argmax(axis=1)
array([2, 0])

还有numpy.argmax模块级别的功能，其功能相同。

>>> np.argmax(a, axis=1)
array([2, 0])