python - Find nth smallest element in numpy array

21
votes

I need to find just the the smallest nth element in a 1D numpy.array.

For example:

a = np.array([90,10,30,40,80,70,20,50,60,0])

I want to get 5th smallest element, so my desired output is 40.

My current solution is this:

result = np.max(np.partition(a, 5)[:5])

However, finding 5 smallest elements and then taking the largest one them seems little clumsy to me. Is there a better way to do it? Am I missing a single function that would achieve my goal?

There are questions with similar titles to this one, but I did not see anything that answered my question.

Edit:

I should've mentioned it originally, but performance is very important for me; therefore, heapq solution though nice would not work for me.

import numpy as np
import heapq

def find_nth_smallest_old_way(a, n):
    return np.max(np.partition(a, n)[:n])

# Solution suggested by Jaime and HYRY    
def find_nth_smallest_proper_way(a, n):
    return np.partition(a, n-1)[n-1]

def find_nth_smallest_heapq(a, n):
    return heapq.nsmallest(n, a)[-1]
#    
n_iterations = 10000

a = np.arange(1000)
np.random.shuffle(a)

t1 = timeit('find_nth_smallest_old_way(a, 100)', 'from __main__ import find_nth_smallest_old_way, a', number = n_iterations)
print 'time taken using partition old_way: {}'.format(t1)    
t2 = timeit('find_nth_smallest_proper_way(a, 100)', 'from __main__ import find_nth_smallest_proper_way, a', number = n_iterations)
print 'time taken using partition proper way: {}'.format(t2) 
t3 = timeit('find_nth_smallest_heapq(a, 100)', 'from __main__ import find_nth_smallest_heapq, a', number = n_iterations)  
print 'time taken using heapq : {}'.format(t3)

Result:

time taken using partition old_way: 0.255564928055
time taken using partition proper way: 0.129678010941
time taken using heapq : 7.81094002724
3
pythonnumpy
Also, might be beneficial to check out docs.python.org/2/library/heapq.html - C.B.
@C.B. the above question is significantly different from mine; it asks for both min and max, and it is for 2D matrix - Akavall
How is this a duplicate? The title sounds similar, but the question itself is very different. Sometimes different questions lead to same answers, but here the answers are also very different. And there is no way an answer in that question is an answer to my question. - Akavall

3 Answers

36
votes

Unless I am missing something, what you want to do is:

>>> a = np.array([90,10,30,40,80,70,20,50,60,0])
>>> np.partition(a, 4)[4]
40

np.partition(a, k) will place the k-th smallest element of a at a[k], smaller values in a[:k] and larger values in a[k+1:]. The only thing to be aware of is that, because of the 0 indexing, the fifth element is at index 4.

5
votes

You can use heapq.nsmallest:

>>> import numpy as np
>>> import heapq
>>> 
>>> a = np.array([90,10,30,40,80,70,20,50,60,0])
>>> heapq.nsmallest(5, a)[-1]
40
1
votes

you don't need call numpy.max():

def nsmall(a, n):
    return np.partition(a, n)[n]