Determine Index of Highest Value in Python's NumPy

Question

I want to generate an array with the index of the highest max value of each row.

a = np.array([ [1,2,3], [6,5,4], [0,1,0] ])
maxIndexArray = getMaxIndexOnEachRow(a)
print maxIndexArray 

[[2], [0], [1]]

There's a np.argmax function but it doesn't appear to do what I want...

Sven Marnach · Accepted Answer · 2010-11-11 00:57:16Z

20

The argmax() function does do what you want:

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

answered Nov 11, 2010 at 0:57

Sven Marnach

607k123 gold badges966 silver badges865 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

This also exists as a top level function, thus the following also works: maxvals = np.argmax(a, axis=1)