Что эквивалентно функции nth_element в Python?

Я хочу реализовать дерево точек выигрыша в 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.

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

@ 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.