python - 在 Python 中进行堆排序以寻求建议

Question

我只需要检索 3 个最小的元素，想知道是否有一种方法可以改进我的以下代码以保持堆大小更小——我认为如果我们只需要将堆大小保持为 3，就足够了。但是在 heapq 中找不到要调整的选项。

换句话说，我想维护一个偶尔更新的三元素堆。

import heapq

def heapsort(iterable):
   h = []
   for value in iterable:
       heapq.heappush(h, value)
   return [heapq.heappop(h) for i in range(len(h))]

if __name__ == "__main__":

   print heapsort([1, 3, 5, 7, 9, 2, 4, 6, 8, 0])

Answer 1

改进代码只获取三个最小元素的方法是将其替换为heapq.nsmallest :

print heapq.nsmallest(3, [1, 3, 5, 7, 9, 2, 4, 6, 8, 0])

输出:

[0, 1, 2]

你可以看看implementation of nsmallest如果您对如何从 heapq 原始函数构建它感到好奇，因为它们正是这样做的。