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

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.

timgeb · Accepted Answer · 2018-12-07 15:35:09Z

2

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])

answered Dec 7, 2018 at 15:35

timgeb

79.2k20 gold badges129 silver badges150 bronze badges

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

2 Comments

Easy with your reputation to close this as a duplicate. You have a gold badge for python

Clearly a duplicate. Should close!