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