python - 如何有效地弹出 heapq 中键最小的所有元素？

Question

我正在做一个模拟实验，并试图让我的代码尽可能高效。一方面，我有一个最小堆优先级队列，我使用heapq模块实现了它。

在整个模拟过程中，我必须用最小的键弹出所有元素。这样做的直接方法是：

    elements = []
    k1, v1 = heapq.heappop(heap)
    elements.append((k1,v1))
    k2, v2 = heap[0] #looks at the smallest item without popping
    while(k1 == k2):
         k1, v1 = heapq.heappop(heap)
         elements.append((k1,v1))
         k2, v2 = heap[0]
    return elements

Answer 1

您展示的一般技术是最有效且最简单的。但是你正在做一些不必要的额外任务。下面是一个小的优化。

elements = []
k1, v1 = heapq.heappop(heap)
elements.append((k1,v1))
while(k1 == heap[0]):
     k2, v2 = heapq.heappop(heap)
     elements.append((k2,v2))
return elements

为了安全起见，您可能应该添加检查以确保您的堆不为空。检查heap[0]堆中何时没有项目将是一件坏事，heapq.heappop如果堆为空则调用。

Answer 2

我打算建议将结构从一堆更改为(k, v)一堆，k然后是一本字典{k:[v]}。这会将您的代码变成：

k = heap[0]
return [(k,v) for v in hash[k]]

和：

hash = defaultdict(list)
heap = []

heappush(heap, (k, v))会成为：

heappush(heap, k)
hash[k].append(v)

heappop(heap)会成为：

k = heappop(heap)
v = hash[k].pop()