Какая эквивалентная функция nth_element в Python?

Я хочу реализовать Vantage Point Tree в Python, но он использует std :: nth_element в C ++.

Поэтому я хочу найти эквивалентную функцию ‘nth_element’ в Python или в numpy.

Обратите внимание, что nth_element будет только частично упорядочивать массив, и это 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);

И теперь вектор может быть:

3,0,2,1,4,5,6,7,9,8

И я не только хочу получить n-й элемент, но также хочу перестроить две части списка, [3,0,2,1,4] и [6,7,9,8].

Кроме того, поддержка nth_element принимает функцию, которая может сравнивать два элемента, например, в приведенном ниже примере вектор представляет собой вектор op DataPoint, а функция DistanceComparator будет сравнивать расстояние между двумя точками с помощью 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,
DistanceComparator(the_v.begin()));

РЕДАКТИРОВАТЬ:

Я использовал ответ bhuvan-venkatesh и написал код для тестирования.

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

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

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

и результат:

2.2217168808
17.0386350155
281.301710844

А потом я сделаю еще тест с использованием кода C ++.

Но есть проблема, при использовании numpy, он всегда будет возвращать новый массив, он будет тратить много памяти, когда мой массив огромен.
Как я могу справиться с этим.
Или я просто должен написать расширение C ++ для Python.

EDIT2:

@ bhuvan-venkatesh Спасибо за рекомендацию функции разделения.

Я использую раздел, как показано ниже:

import numpy

@profile
def for_numpy():
numpy.random.seed(2)
a = numpy.random.rand(1e7)
for i in range(100):
a.partition(numpy.random.randint(1e6))

if __name__ == '__main__':
for_numpy()

и запустил профилировщик как:

python -m memory_profiler profiler_test.py

и результат:

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))

И это не будет копировать весь массив, как:
numpy.partition (a, 3)

Заключение: numpy.ndarray.partition — это то, что я хочу найти.

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)