如何沿一个轴获取numpy数组中最大元素的索引


118

我有一个二维的NumPy数组。我知道如何获取轴上的最大值:

>>> a = array([[1,2,3],[4,3,1]])
>>> amax(a,axis=0)
array([4, 3, 3])

如何获得最大元素的索引?所以我想作为输出array([1,1,0])

Answers:



100
>>> import numpy as np
>>> a = np.array([[1,2,3],[4,3,1]])
>>> i,j = np.unravel_index(a.argmax(), a.shape)
>>> a[i,j]
4

11
请注意,此答案具有误导性。它计算所有方向上数组最大元素的索引,而不是OP要求的沿着给定轴的索引:这是错误的。此外,如果有多个以上最大值,则仅检索第一个最大值的索引:应指出这一点。尝试a = np.array([[1,4,3],[4,3,1]])查看它是否返回i,j==0,1,并忽略了处的解决方案i,j==1,0。对于所有最大值的索引,请改用i,j = where(a==a.max()
gg349

36

argmax()将仅返回每一行的第一个匹配项。 http://docs.scipy.org/doc/numpy/reference/generation/numpy.argmax.html

如果您需要对整形阵列执行此操作,则此方法比unravel

import numpy as np
a = np.array([[1,2,3], [4,3,1]])  # Can be of any shape
indices = np.where(a == a.max())

您还可以更改条件:

indices = np.where(a >= 1.5)

上面以您要求的形式为您提供了结果。另外,您可以通过以下方式将其转换为x,y坐标列表:

x_y_coords =  zip(indices[0], indices[1])

2
这对我不起作用...您是说indices = np.where(a==a.max())第3行吗?
atomh33ls 2014年

你说得对,atomh33ls!感谢您发现这一点。我已经修复了该语句,以包括适当条件的第二个等号。
SevakPrime 2014年

@ SevakPrime,@ atomh33ls指出了第二个错误,.max()而不是.argmax()。请编辑答案
gg349 2015年

@ gg349,这取决于您想要什么。argmax沿轴提供它,这似乎是OP希望其批准eumiro的回答的方式。
SevakPrime 2015年

我看到校正@ atomh33ls和我建议的结果导致数组中最大元素的索引,而OP则沿某个轴询问最大元素。但是请注意,您当前的解决方案导致x_y_coord = [(0, 2), (1, 1)]该结果与@eumiro答案不匹配,并且是错误的。例如,尝试a = array([[7,8,9],[10,11,12]])查看您的代码对此输入没有任何影响。您还提到这比更好unravel,但是@blas发布的解决方案回答了绝对最大值的问题,而不是沿一个轴的问题。
gg349

By using our site, you acknowledge that you have read and understand our Cookie Policy and Privacy Policy.
Licensed under cc by-sa 3.0 with attribution required.