Python中快速获取列表中N个最小或最大元素的方法

Question

我目前有一个使用lambda函数f排序的长列表。然后我从前五个元素中选择一个随机元素。类似于：

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

根据分析器，这是我的程序中的瓶颈。我该如何加快速度？是否有一种快速的内置方法来检索我想要的元素（理论上不需要对整个列表进行排序）？谢谢。

- Sort Me Out Please

---

some_function_of 是昂贵的吗？ - SilentGhost

https://dev59.com/vXE95IYBdhLWcg3wlu-z - fortran

Answer 1

使用heapq.nlargest或heapq.nsmallest。

例如：

import heapq

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

当K相对于N很小时，此方法将以O(N+KlogN)的时间复杂度运行，其中K是返回的元素数量，N是列表大小。相对于正常排序的O(NlogN)，此方法更快。

- interjay

---

嗯。到目前为止，这实际上稍微慢了一点。N是8000，K是5。 - Sort Me Out Please

可能瓶颈在于对某个函数的 N 次调用，而相比之下排序速度更快，这种情况下除了改进该函数外，你无能为力。另一种可能是数据已经接近排序状态，这种情况下 Python 的排序将非常快。 - interjay

你可能是对的。现在会坚持使用heapq.nsmallest，因为它传达了意图。谢谢。 - Sort Me Out Please

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)

- James Hopkin

---

我理解添加到排序函数中的关键参数是用于在内部实现DSU（装饰-排序-去装饰），以便对列表的任何元素仅调用一次可能昂贵的关键函数。我认为您的方法将为相同的列表元素多次调用关键函数。 - PaulMcG