Heaps & Top-K

Top-K, k-way merge, running medians and scheduling.

A heap (priority queue) gives you the smallest (or largest) element in O(1) and inserts or removes in O(log n). Reach for one whenever a problem asks for the top K, the k-th something, or repeatedly needs the next smallest item.

When to use it

Top K / k-th largest — keep a heap of size k
Merging k sorted lists — the heap holds the current head of each list
Running median — two heaps
Scheduling — always process the earliest-ending, highest-priority or cheapest item next
Shortest paths in weighted graphs (Dijkstra)

Python's heapq

heapq is a min-heap on a plain list. For a max-heap, push negated values. Push tuples to order by a key: (priority, tiebreaker, item).

import heapq

h = []
heapq.heappush(h, (3, "c"))
heapq.heappush(h, (1, "a"))
smallest = heapq.heappop(h)      # (1, "a")
heapq.heapify(list_of_items)     # O(n)

Template: top K with a size-k min-heap

import heapq

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

O(n log k) time, O(k) space — better than sorting when k ≪ n, and it works on streams.

Template: k-way merge

import heapq

def merge_k_sorted(lists):
    heap = [(lst[0], i, 0) for i, lst in enumerate(lists) if lst]
    heapq.heapify(heap)
    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

Template: two heaps for a running median

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 at the top of one or both.

import heapq

class MedianFinder:
    def __init__(self):
        self.low, self.high = [], []   # low is a max-heap via negation

    def add(self, x):
        heapq.heappush(self.low, -x)
        heapq.heappush(self.high, -heapq.heappop(self.low))
        if len(self.high) > len(self.low):
            heapq.heappush(self.low, -heapq.heappop(self.high))

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

[!TIP] Practice Top K Frequent Elements · Find Median from Data Stream · Minimum Meeting Rooms