Heaps, priority queues and top-k problems

Know the heap operations and their costs, use a bounded heap for top-k, merge sorted streams, and track a running median.

Heap operations and cost

import heapq
from typing import Iterable, Sequence

# heapq is a MIN-heap. For a max-heap, store negated values.
nums = [5, 1, 9, 3]
heapq.heapify(nums)                  # O(n), in place. Not O(n log n).
print(nums[0])                       # the smallest, O(1)

heapq.heappush(nums, 4)              # O(log n)
smallest = heapq.heappop(nums)       # O(log n)

# push then pop in one call: fewer operations, and the heap never grows
heapq.heappushpop(nums, 0)           # O(log n)

# the k smallest in sorted order, O(n log k)
print(heapq.nsmallest(2, nums))

# a fixed-length heap of the LARGEST k by negating
def k_largest(values: Iterable[int], k: int) -> list[int]:
    if k <= 0:
        return []
    heap: list[int] = []
    for v in values:
        if len(heap) < k:
            heapq.heappush(heap, v)
        elif v > heap[0]:
            heapq.heapreplace(heap, v)     # replace the smallest kept value
    return sorted(heap, reverse=True)

# a priority queue with a tie-breaker, so equal priorities never compare the
# payload. push (priority, sequence, item) and the sequence breaks the tie.
import itertools
def priority_demo() -> list[str]:
    counter = itertools.count()
    pq: list[tuple[int, int, str]] = []
    for task, priority in [("a", 2), ("b", 1), ("c", 1)]:
        heapq.heappush(pq, (priority, next(counter), task))
    return [heapq.heappop(pq)[2] for _ in range(len(pq))]


if __name__ == "__main__":
    print(k_largest([3, 1, 4, 1, 5, 9, 2, 6], 3))    # [9, 6, 5]
    print(priority_demo())                            # ['b', 'c', 'a']

• heapify is O(n) because most nodes are near the bottom and sift down only a little. Calling heappush n times is O(n log n) for the same result.
• A heap is not sorted. Only the root is guaranteed to be the minimum; the rest of the array is a partial order.
• Python's heapq compares tuples element by element. If two priorities are equal it compares the payload, which raises a TypeError for non-comparable items. Always include a unique tie-breaker.
• To change a priority in place, push a new entry with a marker and skip the stale one when it surfaces. Removing from the middle is O(n).

Merging and streaming

import heapq
from typing import Iterator, Sequence

def merge_sorted(lists: Sequence[Sequence[int]]) -> list[int]:
    """k-way merge with a heap of the current head of each list."""
    heap: list[tuple[int, int, int]] = []     # value, list index, position
    for i, lst in enumerate(lists):
        if lst:
            heap.append((lst[0], i, 0))
    heapq.heapify(heap)

    out: list[int] = []
    while heap:
        value, i, j = heapq.heappop(heap)
        out.append(value)
        if j + 1 < len(lists[i]):
            heapq.heappush(heap, (lists[i][j + 1], i, j + 1))
    return out


def median_stream(values: Iterator[int]) -> list[float]:
    """Two heaps: the lower half as a max-heap, the upper half as a min-heap."""
    low: list[int] = []      # negated, so the root is the largest of the lower half
    high: list[int] = []     # the smallest of the upper half
    out: list[float] = []

    for v in values:
        heapq.heappush(low, -v)
        heapq.heappush(high, -heapq.heappop(low))     # move the largest of low up

        if len(high) > len(low):                      # keep low the same size or one bigger
            heapq.heappush(low, -heapq.heappop(high))

        out.append(-low[0] if len(low) > len(high) else (-low[0] + high[0]) / 2)
    return out


def top_k_frequent(words: Sequence[str], k: int) -> list[str]:
    """Count, then keep a bounded heap of size k instead of sorting everything."""
    from collections import Counter
    counts = Counter(words)
    return heapq.nlargest(k, counts, key=counts.get)


if __name__ == "__main__":
    print(merge_sorted([[1, 4, 7], [2, 5, 8], [3, 6, 9]]))
    print(median_stream(iter([5, 15, 1, 3, 8])))   # [5, 10, 5, 4, 5]
    print(top_k_frequent("a b a c a b d".split(), 2))

The median-of-a-stream problem is the clearest demonstration of why two heaps beat one sorted array: insertion stays O(log n) and the median is O(1), whereas re-sorting on every arrival is O(n log n) per element.

Choosing the right technique for top-k

import heapq
import random
from typing import Sequence

def top_k_sort(xs: Sequence[int], k: int) -> list[int]:
    """Sort everything. O(n log n), and O(n) extra space for the copy."""
    return sorted(xs, reverse=True)[:k]

def top_k_heap(xs: Sequence[int], k: int) -> list[int]:
    """Bounded heap. O(n log k) time, O(k) extra space."""
    if k <= 0:
        return []
    heap = list(xs[:k])
    heapq.heapify(heap)
    for v in xs[k:]:
        if v > heap[0]:
            heapq.heapreplace(heap, v)
    return sorted(heap, reverse=True)

def top_k_quickselect(xs: Sequence[int], k: int) -> list[int]:
    """Average O(n) time, but it mutates a copy and has an O(n squared) worst case."""
    data = list(xs)

    def partition(lo: int, hi: int) -> int:
        pivot_index = random.randrange(lo, hi + 1)     # randomised: no adversarial pivot
        data[lo], data[pivot_index] = data[pivot_index], data[lo]
        pivot = data[lo]
        i = lo
        for j in range(lo + 1, hi + 1):
            if data[j] > pivot:                        # descending order
                i += 1
                data[i], data[j] = data[j], data[i]
        data[lo], data[i] = data[i], data[lo]
        return i

    lo, hi = 0, len(data) - 1
    while lo < hi:
        p = partition(lo, hi)
        if p == k - 1:
            break
        if p < k - 1:
            lo = p + 1
        else:
            hi = p - 1
    return sorted(data[:k], reverse=True)

FAQ

Related

Sorting and searching Greedy algorithms and intervals

Last refreshed 2026-09-18.