Быстрый способ получить N Min или Max элементов из списка в Python

Question

В настоящее время у меня есть длинный список, который сортируется с помощью лямбда-функции f.Затем я выбираю случайный элемент из первых пяти элементов.Что-то вроде:

f = lambda x: some_function_of(x, local_variable)
my_list.sort(key=f)
foo = choice(my_list[:4])

Это узкое место в моей программе, согласно профилировщику.Как я могу ускорить процесс?Существует ли быстрый, встроенный способ извлечения нужных мне элементов (теоретически не нужно сортировать весь список).Спасибо.

Answer 1

Используйте heapq.nlargest или heapq.nsmallest.

Например:

import heapq

elements = heapq.nsmallest(4, my_list, key=f)
foo = choice(elements)

Это займет O (N + KlogN) время (где K - количество возвращаемых элементов, а N - размер списка), что быстрее, чем O (NlogN) для нормальной сортировки, когда K мало по отношению к N.

Answer 2

На самом деле это возможно в среднем за линейное время (O (N)).

Вам нужен алгоритм разбиения:

def partition(seq, pred, start=0, end=-1):
    if end == -1: end = len(seq)
    while True:
        while True:
            if start == end: return start
            if not pred(seq[start]): break
            start += 1
        while True:
            if pred(seq[end-1]): break
            end -= 1
            if start == end: return start
        seq[start], seq[end-1] = seq[end-1], seq[start]
        start += 1
        end -= 1

, который может использоваться алгоритмом nth_element:

def nth_element(seq_in, n, key=lambda x:x):
    start, end = 0, len(seq_in)
    seq = [(x, key(x)) for x in seq_in]

    def partition_pred(x): return x[1] < seq[end-1][1]

    while start != end:
        pivot = (end + start) // 2
        seq[pivot], seq[end - 1] = seq[end - 1], seq[pivot]
        pivot = partition(seq, partition_pred, start, end)
        seq[pivot], seq[end - 1] = seq[end - 1], seq[pivot]
        if pivot == n: break
        if pivot < n: start = pivot + 1
        else: end = pivot

    seq_in[:] = (x for x, k in seq)

Учитывая это, просто замените вторую (сортировку) строку на:

nth_element(my_list, 4, key=f)