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.

Glazing

Answer 1

>>> a.argmax(axis=0)

array([1, 1, 0])

Vanbuskirk

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

@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. – Mazurek

Answer 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

Overcrop

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(). – Ascogonium

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

Escobedo

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

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

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

@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. – Escobedo

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 @Vanbuskirk 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. – Ascogonium

@Ascogonium 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. – Escobedo

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

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

Satiety

Answer 5

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

Victorvictoria

Answer 6

Allow me to give an up-to-date answer. We use the function argmax() with the parameter axis that can be defined as follows:

• axis=0, the index of the max element per column will be returned.
• axis=1 the index of the max element per row will be returned.

The following code will help you to better understand.

import numpy as np 
a = np.array([[1,2,3],[4,3,1]])
#This will print the index of the max value for each column.
print(a.argmax(axis=0)) # output: [1 1 0]
#This will print the index of the max value for each row.
print(a.argmax(axis=1)) # output: [2 0]

Lugsail