Heaps and Priority Queues

A heap gives you the smallest (or largest) item in $O(1)$ and lets you insert and remove in $O(\log n)$. That makes it the right tool whenever a problem keeps asking "what is the best remaining item?" while the set of items changes: the $k$ largest elements, the next task to run, the lightest edge to explore, the running median. The key patterns are top $k$ with a heap of size $k$, merging sorted streams, and two heaps for medians.

1. When this pattern applies

• You need the $k$ largest or smallest items, or the $k$th one.
• You repeatedly need the minimum or maximum of a collection that changes.
• You merge several sorted sequences.
• You process events in priority or time order (scheduling, simulations, Dijkstra's algorithm).
• The data arrives as a stream and you cannot sort it all.

2. Heaps in Python

Python's heapq implements a min-heap on a list.

import heapq

h = []
for x in [5, 1, 8, 3]:
    heapq.heappush(h, x)               # O(log n)
assert h[0] == 1                       # the smallest is always at index 0, O(1)
assert heapq.heappop(h) == 1           # O(log n)
assert heapq.heappop(h) == 3

nums = [9, 4, 7, 1]
heapq.heapify(nums)                    # O(n): turn a list into a heap in place
assert nums[0] == 1

assert heapq.nlargest(2, [5, 1, 8, 3]) == [8, 5]
assert heapq.nsmallest(2, [5, 1, 8, 3]) == [1, 3]

A max-heap is simulated by pushing negated values.

mx = []
for x in [5, 1, 8, 3]:
    heapq.heappush(mx, -x)
assert -mx[0] == 8

Heap entries can be tuples, compared left to right, so you can attach priorities or payloads: (priority, item). If priorities tie and the items are not comparable, add a tie-breaker counter: (priority, counter, item).

How a heap works, in outline: it is a complete binary tree stored in an array where each parent is at most its children. Insertion adds at the end and sifts up. Removing the minimum swaps the last element to the root and sifts down. Both take $O(\log n)$ because the tree has $\log n$ levels. Be ready to implement it:

class MinHeap:
    def __init__(self):
        self.a = []

    def push(self, x):
        self.a.append(x)
        i = len(self.a) - 1
        while i > 0 and self.a[(i - 1) // 2] > self.a[i]:       # sift up
            self.a[(i - 1) // 2], self.a[i] = self.a[i], self.a[(i - 1) // 2]
            i = (i - 1) // 2

    def pop(self):
        top = self.a[0]
        last = self.a.pop()
        if self.a:
            self.a[0] = last
            i, n = 0, len(self.a)
            while True:                                          # sift down
                smallest = i
                for c in (2 * i + 1, 2 * i + 2):
                    if c < n and self.a[c] < self.a[smallest]:
                        smallest = c
                if smallest == i:
                    break
                self.a[i], self.a[smallest] = self.a[smallest], self.a[i]
                i = smallest
        return top

mh = MinHeap()
for x in [5, 3, 8, 1, 9, 2]:
    mh.push(x)
assert [mh.pop() for _ in range(6)] == [1, 2, 3, 5, 8, 9]

3. Pattern A: top $k$ with a heap of size $k$

To keep the $k$ largest items, maintain a min-heap of size $k$. For each new item, push it, and if the heap exceeds $k$, pop the smallest. At the end, the heap holds the $k$ largest. The smallest of them is at the root, which is also the $k$th largest.

The heap is small ($k$ items), so each operation is $O(\log k)$, for $O(n\log k)$ total, better than sorting when $k\ll n$.

def kth_largest(nums, k):
    heap = []
    for x in nums:
        heapq.heappush(heap, x)
        if len(heap) > k:
            heapq.heappop(heap)           # drop the smallest of the k+1 candidates
    return heap[0]

assert kth_largest([3, 2, 1, 5, 6, 4], 2) == 5
assert kth_largest([3, 2, 3, 1, 2, 4, 5, 5, 6], 4) == 4

A stream version keeps the heap between calls:

class KthLargest:
    def __init__(self, k, nums):
        self.k = k
        self.heap = []
        for x in nums:
            self.add(x)

    def add(self, val):
        heapq.heappush(self.heap, val)
        if len(self.heap) > self.k:
            heapq.heappop(self.heap)
        return self.heap[0]

kl = KthLargest(3, [4, 5, 8, 2])
assert kl.add(3) == 4 and kl.add(5) == 5 and kl.add(10) == 5 and kl.add(9) == 8

Top $k$ frequent elements with a heap of (count, value):

from collections import Counter

def top_k_frequent(nums, k):
    count = Counter(nums)
    return [v for _, v in heapq.nlargest(k, ((c, v) for v, c in count.items()))]

assert sorted(top_k_frequent([1, 1, 1, 2, 2, 3], 2)) == [1, 2]

$k$ closest points to the origin. Use a max-heap of size $k$ on distance, keeping the $k$ smallest distances (negate the distance).

def k_closest(points, k):
    heap = []                                       # max-heap via negated distance
    for x, y in points:
        d = x * x + y * y                           # squared distance: no sqrt needed
        heapq.heappush(heap, (-d, x, y))
        if len(heap) > k:
            heapq.heappop(heap)
    return [[x, y] for _, x, y in heap]

assert sorted(k_closest([[1, 3], [-2, 2]], 1)) == [[-2, 2]]
assert sorted(k_closest([[3, 3], [5, -1], [-2, 4]], 2)) == [[-2, 4], [3, 3]]

Note the direction: to keep the $k$ largest, use a min-heap, and to keep the $k$ smallest, use a max-heap. That "opposite heap" idea trips people up, so say it out loud.

4. Pattern B: merge $k$ sorted sequences

Keep a heap holding the current front element of each list. Repeatedly pop the smallest, then push the next element from the same list. The heap never exceeds $k$ items, so the cost is $O(N\log k)$ for $N$ total elements.

def merge_k_sorted(lists):
    heap = []
    for i, lst in enumerate(lists):
        if lst:
            heapq.heappush(heap, (lst[0], i, 0))      # value, which list, index in it
    out = []
    while heap:
        val, i, j = heapq.heappop(heap)
        out.append(val)
        if j + 1 < len(lists[i]):
            heapq.heappush(heap, (lists[i][j + 1], i, j + 1))
    return out

assert merge_k_sorted([[1, 4, 5], [1, 3, 4], [2, 6]]) == [1, 1, 2, 3, 4, 4, 5, 6]
assert merge_k_sorted([[], [1]]) == [1]
assert merge_k_sorted([]) == []

The linked-list version of this problem uses the same idea with node references. Because nodes are not comparable in Python, include an index as a tie-breaker:

class ListNode:
    def __init__(self, val=0, next=None):
        self.val, self.next = val, next

def merge_k_lists(lists):
    heap = [(node.val, i, node) for i, node in enumerate(lists) if node]
    heapq.heapify(heap)
    dummy = tail = ListNode()
    while heap:
        _, i, node = heapq.heappop(heap)
        tail.next = node
        tail = node
        if node.next:
            heapq.heappush(heap, (node.next.val, i, node.next))
    return dummy.next

def make(vals):
    d = t = ListNode()
    for v in vals:
        t.next = ListNode(v); t = t.next
    return d.next

out, node = [], merge_k_lists([make([1, 4, 5]), make([1, 3, 4]), make([2, 6])])
while node:
    out.append(node.val); node = node.next
assert out == [1, 1, 2, 3, 4, 4, 5, 6]

An alternative, divide and conquer, merges lists pairwise and also costs $O(N\log k)$ with $O(1)$ extra space on linked lists. Mention both.

5. Pattern C: two heaps for the running median

To find the median of a data stream, keep the lower half in a max-heap and the upper half in a min-heap, balanced so their sizes differ by at most one. The median is the top of the larger heap, or the average of both tops.

class MedianFinder:
    def __init__(self):
        self.low = []        # max-heap (negated): the smaller half
        self.high = []       # min-heap: the larger half

    def add_num(self, num):
        heapq.heappush(self.low, -num)
        heapq.heappush(self.high, -heapq.heappop(self.low))   # move the largest of low to high
        if len(self.high) > len(self.low):                      # rebalance: low may hold one extra
            heapq.heappush(self.low, -heapq.heappop(self.high))

    def find_median(self):
        if len(self.low) > len(self.high):
            return -self.low[0]
        return (-self.low[0] + self.high[0]) / 2

m = MedianFinder()
m.add_num(1); m.add_num(2)
assert m.find_median() == 1.5
m.add_num(3)
assert m.find_median() == 2
for x in [10, 20, 30]:
    m.add_num(x)
assert m.find_median() == 6.5

Time $O(\log n)$ per insertion, $O(1)$ per median query. The invariant is the whole solution: every element of low is at most every element of high, and the sizes are balanced.

A harder variant, the sliding window median, needs deletion of arbitrary elements. A heap with lazy deletion, or a sorted structure, handles it.

6. Pattern D: scheduling and greedy with a heap

Task scheduler / reorganise string. When you must repeatedly pick the most frequent remaining item subject to a constraint, keep counts in a max-heap.

Reorganise string (no two adjacent characters equal): repeatedly place the most frequent character that is not the one just placed.

from collections import Counter

def reorganize_string(s):
    heap = [(-c, ch) for ch, c in Counter(s).items()]
    heapq.heapify(heap)
    result = []
    prev = (0, "")                                    # the held-back previous character
    while heap:
        count, ch = heapq.heappop(heap)
        result.append(ch)
        if prev[0] < 0:
            heapq.heappush(heap, prev)               # the previous one becomes available again
        prev = (count + 1, ch)                       # one fewer of this character remains
    out = "".join(result)
    return out if len(out) == len(s) else ""

assert reorganize_string("aab") in ("aba",)
assert reorganize_string("aaab") == ""
r = reorganize_string("vvvlo")
assert all(r[i] != r[i + 1] for i in range(len(r) - 1)) and sorted(r) == sorted("vvvlo")

Meeting rooms II with a heap of end times is in the sorting chapter. Dijkstra's shortest path uses a heap of (distance, node) and is in the graphs chapter. Huffman coding, connecting ropes at minimum cost and IPO / maximise capital are the same idea: always combine or choose the cheapest or best available item.

def min_cost_to_connect_ropes(ropes):
    heapq.heapify(ropes)
    cost = 0
    while len(ropes) > 1:
        a, b = heapq.heappop(ropes), heapq.heappop(ropes)
        cost += a + b
        heapq.heappush(ropes, a + b)
    return cost

assert min_cost_to_connect_ropes([4, 3, 2, 6]) == 29
assert min_cost_to_connect_ropes([5]) == 0

7. Choosing between a heap, sorting and quickselect

A heap does not support fast search for an arbitrary element, and it does not keep items fully sorted: only the root is guaranteed to be the extreme.

8. Common mistakes

• Using a max-heap where a min-heap of size $k$ is needed (or the reverse) for top $k$.
• Forgetting that heapq is a min-heap, so a max-heap needs negation, and negating also affects tuple ordering.
• Pushing unorderable items (such as nodes) without a tie-breaker, which raises an error when priorities tie.
• Reading heap[0] on an empty heap. Guard against it.
• Assuming the list is sorted because it is a heap. Only heap[0] is guaranteed.
• Using sorted repeatedly on a growing list instead of a heap, giving $O(n^2\log n)$.
• Breaking the two-heap invariant by not rebalancing, or by pushing into the wrong heap first.

9. Practice set

1. Kth largest element in an array, kth largest in a stream.
2. Top $k$ frequent elements, top $k$ frequent words (with tie-breaking), sort characters by frequency.
3. $k$ closest points to origin, $k$ closest elements in a sorted array.
4. Merge $k$ sorted lists, smallest range covering elements from $k$ lists.
5. Find median from data stream, sliding window median.
6. Task scheduler, reorganise string, rearrange string $k$ distance apart.
7. Last stone weight, minimum cost to connect sticks.
8. Meeting rooms II, car pooling, the skyline problem.
9. Ugly number II, super ugly number (generate in order with a heap).
10. IPO (maximise capital) and course schedule III (greedy with a heap).