Get the index of the highest value inside a numpy array for each row?

Question

I have a numpy array of 30 rows and 4 columns, for each row I need to get the index where the highest value is located.

So for an array like this

a = np.array([[0, 1, 2],[7, 4, 5]])

I would like to get a list with the indices 2 for the first row and 0 for the second row.

I tried with the numpy function argmax as follows:

for i in range(len(a)):
    results=[np.argmax(a)]
return (results)

but I just get the global maximum, does anyone know how to go around this?

Thank you very much for the help.

Answer 1

Use the argmax method with axis=1 in order to work on the rows.

>>> import numpy as np
>>> a = np.array([[0, 1, 2],[7, 4, 5]])
>>> a.argmax(axis=1)
array([2, 0])

There is also the numpy.argmax module level function which works just the same.

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