heapq

  • A heap is a tree-like data structure in which the child nodes have a sort-order relationship with the parents
  • Use a list or array organized so that the children of element N are at positions 2 N + 1 and 2 N + 2
  • Max-heap, the value of parent is greater than the value of its children
  • Min-heap, the value of parent is less than the value of its children
  • heapq implements a min-heap sort algorithm suitable for use with Python’s lists

Initialization

In [12]:
import math
from io import StringIO


def show_tree(tree, total_width=36, fill=' '):
    """Pretty-print a tree."""
    output = StringIO()
    last_row = -1
    for i, n in enumerate(tree):
        if i:
            row = int(math.floor(math.log(i + 1, 2)))
        else:
            row = 0
        if row != last_row:
            output.write('\n')
        columns = 2 ** row
        col_width = int(math.floor(total_width / columns))
        output.write(str(n).center(col_width, fill))
        last_row = row
    print(output.getvalue())
    print('-' * total_width)
    print()
In [31]:
import heapq

h = []
l = [16, 14, 10, 8, 7, 9, 3, 2, 4, 1]

# heappush, push the value item onto the heap
for i in l:
    heapq.heappush(h, i)
show_tree(h)

l = [16, 14, 10, 8, 7, 9, 3, 2, 4, 1]
heapq.heapify(l) # transform list x into a heap
show_tree(l)

                 1                  
        2                 7         
    4        3        14       9    
 16  8   10 
------------------------------------


                 1                  
        2                 3         
    4        7        9        10   
 8   16  14 
------------------------------------

Methods

In [54]:
l = [16, 14, 10, 8, 7, 9, 3, 2, 4, 1]
heapq.heapify(l)

# heappop, pop and return the smallest item from the heap
e = heapq.heappop(l)
show_tree(l)

# heapq.heappushpop, push item on the heap, then pop and return the smallest item from the heap
e = heapq.heappushpop(l, 5)
show_tree(l)

# heapq.heapreplace, pop and return the smallest item from the heap, and also push the new item
e = heapq.heapreplace(l, 6)
show_tree(l)

# heapq.nlargest
l2 = heapq.nlargest(4, l)
print(l2) # [16, 14, 10, 9]

# heapq.nsmallest
l3 = heapq.nsmallest(4, l)
print(l3) # [4, 5, 6, 7]

                 2                  
        4                 3         
    8        7        9        10   
 14  16 
------------------------------------


                 3                  
        4                 5         
    8        7        9        10   
 14  16 
------------------------------------


                 4                  
        6                 5         
    8        7        9        10   
 14  16 
------------------------------------

[16, 14, 10, 9]
[4, 5, 6, 7]

Heap sort

In [34]:
l = [16, 14, 10, 8, 7, 9, 3, 2, 4, 1]
heapq.heapify(l)

s = [heapq.heappop(l) for i in range(len(l))]
print(s) # [1, 2, 3, 4, 7, 8, 9, 10, 14, 16]

[1, 2, 3, 4, 7, 8, 9, 10, 14, 16]

Priority Queue

In [42]:
l = [16, 14, 10, 8, 7, 9, 3, 2, 4, 1]
heapq.heapify(l)
show_tree(l)

# push element into queue
heapq.heappush(l, 5)
show_tree(l)

# pop element
e = heapq.heappop(l)
show_tree(l)

                 1                  
        2                 3         
    4        7        9        10   
 8   16  14 
------------------------------------


                 1                  
        2                 3         
    4        5        9        10   
 8   16  14  7  
------------------------------------


                 2                  
        4                 3         
    7        5        9        10   
 8   16  14 
------------------------------------

Reference