首页 > 解决方案 > 获取每行的numpy数组中最大值的索引?

问题描述

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

所以对于这样的数组

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)

但我只是得到了全球最大值,有谁知道如何解决这个问题?

非常感谢你的帮助。

标签: pythonarraysnumpy

解决方案


使用argmaxwith 方法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])

推荐阅读