Clever Geek Handbook

Following argmax values ​​in python

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

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

However, I want to get the next largest values, is there a way to remove the getMax values ​​from the original array and then execute argmax again?

+3


source to share


2 answers


Use the command np.argsort(a, axis=-1, kind='quicksort', order=None)

, but with the appropriate choice of arguments (below).

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

The default order is small to large. So sort with -dist

(for quick coding). Caveat: Execution -dist

results in a new array being created, which might bother you if it's dist

huge. See the bottom of the post for a better alternative.

Here's 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 nth entry in the line L

gives the location of the nth largest element x

.

x

does not change.



L[:,0]

will give the same result 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 the hypotheticalargsecondmax(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])
+5


source


If speed is important to you, it is useful to use argpartition rather than argsort.

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

import numpy as np 

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

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


Function %timeit

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

.

Hope this is helpful.

0


source






More articles:


All Articles