Heaps & Priority Queues — Top K, Kth Largest and Running Median

A heap (priority queue) always gives you the smallest (min-heap) or largest (max-heap) item in O(1), and adds or removes items in O(log n). Use it whenever a question keeps asking "what is the current smallest/largest?" while items keep arriving.

Operation	Cost
Peek min/max	O(1)
Push / pop	O(log n)
Build a heap from n items	O(n)

Using heaps in Python and Java

import heapq

h = []
heapq.heappush(h, 5)
heapq.heappush(h, 1)
heapq.heappush(h, 3)
heapq.heappop(h)        # 1 — Python's heapq is a MIN-heap

# Max-heap trick: push negatives
heapq.heappush(h, -10)  # then negate again when you pop

PriorityQueue<Integer> minHeap = new PriorityQueue<>();
PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder());
minHeap.offer(5); minHeap.offer(1);
minHeap.poll();   // 1

Under the hood a heap is an array arranged as a complete binary tree where every parent is ≤ its children (min-heap). You rarely need to implement it — but know the costs.

Pattern 1 — Kth largest: keep a heap of size K

To find the k-th largest element, keep a min-heap of the k largest seen so far. Its top is the smallest of those k — exactly the k-th largest.

def find_kth_largest(nums, k):
    heap = []
    for x in nums:
        heapq.heappush(heap, x)
        if len(heap) > k:
            heapq.heappop(heap)      # drop the smallest — it can't be in the top k
    return heap[0]

O(n log k) time, O(k) space — better than sorting (O(n log n)) when k is small.

Counter-intuitive but important: for the k largest, use a min-heap (to evict small ones). For the k smallest, use a max-heap.

Pattern 2 — Top K frequent

Count with a hash map, then keep a size-k heap ordered by frequency:

from collections import Counter

def top_k_frequent(nums, k):
    count = Counter(nums)
    return [x for x, _ in heapq.nlargest(k, count.items(), key=lambda kv: kv[1])]

int[] topKFrequent(int[] nums, int k) {
    Map<Integer, Integer> count = new HashMap<>();
    for (int x : nums) count.merge(x, 1, Integer::sum);
    PriorityQueue<int[]> heap = new PriorityQueue<>((a, b) -> a[1] - b[1]); // min by frequency
    for (var e : count.entrySet()) {
        heap.offer(new int[]{e.getKey(), e.getValue()});
        if (heap.size() > k) heap.poll();
    }
    int[] res = new int[k];
    for (int i = 0; i < k; i++) res[i] = heap.poll()[0];
    return res;
}

K Closest Points to Origin is the same idea, comparing by distance.

Pattern 3 — Merge K sorted lists

Put the head of each list in a min-heap. Repeatedly pop the smallest, append it to the answer, and push that node's next. The heap never holds more than k items → O(N log k) for N total nodes.

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)        # i breaks ties (nodes aren't comparable)
        tail.next = tail = node
        if node.next:
            heapq.heappush(heap, (node.next.val, i, node.next))
    return dummy.next

Pattern 4 — Two heaps for the running median

Find Median from Data Stream: keep the smaller half in a max-heap and the larger half in a min-heap, balanced in size. The median is at the tops.

class MedianFinder:
    def __init__(self):
        self.low = []    # max-heap (store negatives)
        self.high = []   # min-heap
    def addNum(self, x):
        heapq.heappush(self.low, -x)
        heapq.heappush(self.high, -heapq.heappop(self.low))   # move the biggest of low to high
        if len(self.high) > len(self.low):
            heapq.heappush(self.low, -heapq.heappop(self.high))
    def findMedian(self):
        if len(self.low) > len(self.high):
            return -self.low[0]
        return (-self.low[0] + self.high[0]) / 2

Each insertion is O(log n), each median query O(1).

Pattern 5 — Greedy scheduling with a heap

Problems like Last Stone Weight (repeatedly smash the two heaviest) and Task Scheduler repeatedly need the current largest — a max-heap makes each step O(log n).

How to recognise a heap problem

• "K-th largest/smallest", "top K", "K closest".
• "Merge K sorted …".
• "Running/streaming median", "repeatedly take the largest/smallest".

Common mistakes

1. Using a max-heap for the k largest (it works, but costs O(n log n) memory/time — use a size-k min-heap).
2. Forgetting Python's heapq is min-only — negate values for a max-heap.
3. Pushing objects that can't be compared (Python): add a tie-breaker index to the tuple.
4. Sorting everything when only the top k are needed.

Key takeaways

• Heap = fast access to the current min/max: peek O(1), push/pop O(log n).
• Size-k heap → Top K / Kth in O(n log k).
• Two heaps → running median. Heap of list heads → merge K sorted lists.