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...

Answer 1

The argmax() function does do what you want:

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