堆在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]

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