Find nth smallest element in numpy array

Question

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

Asked By: Akavall

||

Source

Answer 1

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

Answered By: arshajii

Answer 2

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

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

Answered By: HYRY

Answer 3

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+1-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.

Answered By: Jaime

Find nth smallest element in numpy array

Question:

Answers: