从Python中的列表中获取N Min或Max元素的快速方法

Question

我目前有一个很长的列表，正在使用lambda函数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是列表大小），当K相对于N小时，其比正常排序的O（NlogN）快。

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)