获取每行的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])