How to get the index of a maximum element in a NumPy array along one axis

Question

I have a 2 dimensional NumPy array. I know how to get the maximum values over axes:

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

How can I get the indices of the maximum elements? I would like as output array([1,1,0]) instead.

Solution 1

>>> a.argmax(axis=0)

array([1, 1, 0])

Solution 2

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

Solution 3

argmax() will only return the first occurrence for each row. http://docs.scipy.org/doc/numpy/reference/generated/numpy.argmax.html

If you ever need to do this for a shaped array, this works better than 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())

You can also change your conditions:

indices = np.where(a >= 1.5)

The above gives you results in the form that you asked for. Alternatively, you can convert to a list of x,y coordinates by:

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

Solution 4

There is argmin() and argmax() provided by numpy that returns the index of the min and max of a numpy array respectively.

Say e.g for 1-D array you'll do something like this

import numpy as np

a = np.array([50,1,0,2])

print(a.argmax()) # returns 0
print(a.argmin()) # returns 2

And similarly for multi-dimensional array

import numpy as np

a = np.array([[0,2,3],[4,30,1]])

print(a.argmax()) # returns 4
print(a.argmin()) # returns 0

Note that these will only return the index of the first occurrence.

Solution 5

v = alli.max()
index = alli.argmax()
x, y = index/8, index%8

View more solutions

Comments

[Lee]

This didn't work for me... Do you mean indices = np.where(a==a.max()) in line 3?

[SevakPrime]

You are right, atomh33ls! Thanks for spotting that. I've fixed that statement to include the second equals sign for the proper conditional.

[gg349]

@SevakPrime, there was a second error pointed out by @atomh33ls, .max() instead of .argmax(). Please edit the answer

[SevakPrime]

@gg349, it depends on what you want. argmax provides it along an axis which seems to be the way the OP wants it having approved that answer by eumiro.

[gg349]

I see that the correction @atomh33ls and I propose leads to the index of the largest element(s) of the array, while the OP was asking about the largest elements along a certain axis. Notice however that your current solution leads to x_y_coord = [(0, 2), (1, 1)] that does NOT match @eumiro answer, and is wrong. For example, try with a = array([[7,8,9],[10,11,12]]) to see that your code does not have any hit on this input. You also mention that this works better than unravel, but the solution posted by @blas answer the problem of the absolute maximum, not jsut along one axis.

[gg349]

Notice that this answer is misleading. It calculates the index of the maximum element of the array across all axis, not along a given axis as the OP asks: it is wrong. Moreover, if there is more than one maximum, it retrieves the indices of only the first maximum: this should be pointed out. Try with a = np.array([[1,4,3],[4,3,1]]) to see that it returns i,j==0,1, and neglects the solution at i,j==1,0. For the indices of all the maxima use instead i,j = where(a==a.max().

[SevakPrime]

@gg349 and atomh33ls, yes you're correct. This had worked for me as posted originally, but it clearly doesn't work now. Not sure if it has anything to do with any updates. Thanks for your corrections. I've made the edit.

[Løiten]

This worked fine with me. However, notice that if one is to compare reals, one should use np.where(np.isclose(a, a.max())) due to finite precision

[Priyom saha]

this works fine for integers but what can I do for float values and the numbers between 0 and 1

[Philippe Carphin]

@Priyom saha This works for an array of floats, the resulting array is an array of indices where the largest floats are in each column. In the first column, the second element is the largest, in the second column the second element is the largest, and in the third column, the first element is the largest.