Next argmax values in python

Question

I have a function that returns the argmax from a large 2d array

getMax = np.argmax(dist, axis=1)

However I want to get the next biggest values, is there a way of removing the getMax values from the original array and then performing argmax again?

Answer 1

Use the command np.argsort(a, axis=-1, kind='quicksort', order=None), but with appropriate choice of arguments (below).

here is the documentation. Note "It returns an array of indices of the same shape as a that index data along the given axis in sorted order."

The default order is small to large. So sort with -dist (for quick coding). Caution: doing -dist causes a new array to be generated which you may care about if dist is huge. See bottom of post for a better alternative there.

Here is an example:

x = np.array([[1,2,5,0],[5,7,2,3]])
L = np.argsort(-x, axis=1)

print L
[[2 1 0 3]
 [1 0 3 2]]

x  
array([[1, 2, 5, 0],
   [5, 7, 2, 3]])

So the n'th entry in a row of L gives the locations of the n'th largest element of x.

x is unchanged.

L[:,0] will give the same output as np.argmax(x)

L[:,0]
array([2, 1])

np.argmax(x,axis=1)
array([2, 1])

and L[:,1] will give the same as a hypothetical argsecondmax(x)

L[:,1]
array([1, 0])

If you don't want to generate a new list, so you don't want to use -x:

L = np.argsort(x, axis=1)

print L
[[3 0 1 2]
 [2 3 0 1]]

L[:,-1]
array([2, 1])

L[:,-2]
array([1, 0])

Answer 2

If speed is important to you, using argpartition rather than argsort could be useful.

For example, to return the n largest elements from a list:

import numpy as np 

l = np.random.random_integer(0, 100, 1e6)

top_n_1 = l[np.argsort(-l)[0:n]]
top_n_2 = l[np.argpartition(l, -n)[-n:]]

The %timeit function in ipython reports 10 loops, best of 3: 56.9 ms per loop for top_n_1 and 100 loops, best of 3: 8.06 ms per loop for top_n_2.

I hope this is useful.