1

I am working on a simulation experiment and trying to make my code as efficient as possible. In one part, I have a min heap priority queue that I have implemented using heapq module.

Throughout the simulation, I have to pop all elements with the smallest key. The straightforward approach to do this would be:

    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
Vadim Kotov
  • 8,084
  • 8
  • 48
  • 62
Joe
  • 79
  • 3

2 Answers2

1

The general technique you show is the most efficient as well as being straightforward. But you're doing extra assignments that aren't really necessary. Below is a minor optimization.

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

To be on the safe side, you probably should add checks to make sure your heap isn't empty. Checking heap[0] when there are no items in the heap would be a bad thing, as would calling heapq.heappop if the heap is empty.

Jim Mischel
  • 131,090
  • 20
  • 188
  • 351
  • `while k1 == heap[0][0]:` in this case, as each element in `heap` is a `(k, v)` tuple. And simply adding `while heap and k1 == heap[0][0]:` will ensure that you never try to do anything to an empty heap. – Eknoma Oct 10 '22 at 19:42
0

I was going to suggest a change of structure from a heap of (k, v) to a heap of k and a dictionary of {k:[v]}. This would turn your code into:

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

With:

hash = defaultdict(list)
heap = []

heappush(heap, (k, v)) would become:

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

heappop(heap) would become:

k = heappop(heap)
v = hash[k].pop()
Dan D.
  • 73,243
  • 15
  • 104
  • 123
  • How is that more efficient? He's doing the same amount of work. Your code just complicates it by adding another data structure. – Jim Mischel Jan 21 '19 at 18:54