Colin Ji Colin Ji - 2 months ago 3x
C++ Question

What's the equivalent 'nth_element' function in Python?

I want to implement Vantage Point Tree in the python, but it use the std::nth_element in the C++.

So I want to find the equivalent 'nth_element' function in Python or in numpy.

Notice, the nth_element would only partial order an array, and it's O(N).

int the_array[10] = {4,5,7,3,6,0,1,2,9,8};
std::vector<int> the_v(the_array,the_array+10);
std::nth_element (the_v.begin()+0, the_v.begin()+5, the_v.begin()+10);

And now the vector could be:


And I not only want to get the nth element, but also want to get the re-arrange the two part of list, [3,0,2,1,4] and [6,7,9,8].

Moreover, nth_element support accept a function that could compare the two elements, such as, in the below as, the vector is a vector op DataPoint, and the DistanceComparator function will compare the two points distance with the_v.begin():

vector<DataPoint> the_v;
for(int n = 0; n < N; n++) the_v[n] = DataPoint(D, n, X + n * D);
std::nth_element (the_v.begin()+0, the_v.begin()+5, the_v.begin()+10,


I've used the bhuvan-venkatesh's answer, and write some code to test.

partition_timer = timeit.Timer("numpy.partition(a, 10000)",
"import numpy;numpy.random.seed(2);"+
"a = numpy.random.rand(10000000)")

sort_timer = timeit.Timer("numpy.sort(a)",
"import numpy;numpy.random.seed(2);"+
"a = numpy.random.rand(10000000)")

sorted_timer = timeit.Timer("sorted(a)",
"import numpy;numpy.random.seed(2);"+
"a = numpy.random.rand(10000000)")

and the result:


And then, I will do more test using C++ code.

But there is a problem, when using the numpy, it would always return a new array, it will waste a lot of memory, when my array is huge.
How can I handle it.
Or I just have to write a C++ extend for python.


@bhuvan-venkatesh Thanks for recommending the partition function.

I use partition like the below:

import numpy

def for_numpy():
a = numpy.random.rand(1e7)
for i in range(100):

if __name__ == '__main__':

and ran the profiler like:

python -m memory_profiler

and the result is:

Line # Mem usage Increment Line Contents
25 23.613 MiB 0.000 MiB @profile
26 def for_numpy():
27 23.613 MiB 0.000 MiB numpy.random.seed(2)
28 99.934 MiB 76.320 MiB a = numpy.random.rand(1e7)
29 100.004 MiB 0.070 MiB for i in range(100):
30 100.004 MiB 0.000 MiB a.partition(numpy.random.randint(1e6))

And it would not copy the whole array like:
numpy.partition(a, 3)

Conclusion: numpy.ndarray.partition is the one I want to find.


Just make sure that the numpy partition will create two new arrays, meaning that you will quickly make a lot of new arrays. They are more efficient than python lists but will not do the exact same thing as in c++.

If you want the exact element then you can do a filter search which will still be O(n)

array = np.array(...)
partition = np.partition(array, 5) # O(n)
element = np.where(partition==array[5]) # O(n)
left, right = partition[:element], partition[element+1:] # O(n)

So your new code is slower but that is the python-y way to do it.


So you need a comparator? Apart from writing a small function of your own there is no way -- in pure numpy as a keyword -- because each numpy operation is implemented in highly optimized c-code meaning that passing in a python function or a python lambda would force numpy to go to the object level every time and eval.

numpy.vectorize goes to the object level, but in the end you will have to write your own code; Rosetta code has the impelmentation if you want to create a more "optimized algorithm". (I put that in quotes because with python objects you will still be much slower than c or numpy code because of object level access). If speed is your true concern, but you want the python readability consider making an extension with cython.