我只需要检索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])
改进代码以仅获取三个最小元素的方法是将其替换为heapq.nsmallest
:
print heapq.nsmallest(3, [1, 3, 5, 7, 9, 2, 4, 6, 8, 0])
输出:
[0, 1, 2]
如果您对如何从原始函数构建它感到好奇,nsmallest
可以看看的实现heapq
,因为它们正是这样做的。
本文收集自互联网,转载请注明来源。
如有侵权,请联系 [email protected] 删除。
我来说两句