Heap & Top-K

At a glance

The instinct

Sorting gives you everything in order, but you usually only need the best few or the best one right now. A heap is a partially ordered tree that keeps just one promise: the root is the minimum (or maximum). That single promise costs O(log n) per insert or removal and is exactly enough for two families of problems:

• Top-K: keep a heap of size k. To keep the k largest, use a min-heap - its root is the weakest member of the club, the one to evict when someone better arrives.
• Greedy simulation / scheduling: at each moment, push everything that has become available and pop the best one. CPU schedulers, server assignment, "use ladders on the tallest climbs".

Recognise it

Not this pattern if...

• You need arbitrary deletion or floor/ceiling of a value - use an ordered set (TreeMap).
• k is close to n, or you need the whole order - just sort.
• Values are small bounded integers (frequencies <= n) - bucket sort gives O(n) top-K without a heap.
• You need one kth element on a static array as fast as possible - quickselect averages O(n).

Memory hack

Analogy - a hospital emergency room. Patients arrive at different times (sort by arrival), and the triage nurse always calls the most urgent waiting patient next (poll the heap), not whoever came first. When nobody is waiting, the clock jumps to the next arrival.

Say it out loud in the interview:

1. "I only need the k best, so I keep a heap of size k - min-heap for largest, max-heap for smallest."
2. "My comparator is Integer.compare / Comparator.comparingInt, never subtraction."
3. "Each element costs O(log k), so overall O(n log k), and the heap holds at most k items."

How it works

Trace of Kth Largest Element in a Stream with k = 3, stream 4, 5, 8, 2, 3, 10, 9:

For scheduling problems the loop shape is different:

flowchart TD
  A["sort tasks by arrival"] --> B{"tasks left or heap not empty?"}
  B -->|no| Z["done"]
  B -->|yes| C{"heap empty and next task in future?"}
  C -->|yes| D["jump time to next arrival"]
  C -->|no| E["push every task with arrival <= time"]
  D --> E
  E --> F["pop best task, run it, advance time"]
  F --> B

Java template

import java.util.*;

class TopK {
    // 1) Top-K by a score: keep a MIN-heap of the k best (largest scores).
    static List<Integer> topKFrequent(int[] nums, int k) {
        Map<Integer, Integer> freq = new HashMap<>();
        for (int x : nums) freq.merge(x, 1, Integer::sum);

        PriorityQueue<Integer> heap =                 // root = weakest of the current top k
            new PriorityQueue<>(Comparator.comparingInt(freq::get));
        for (int key : freq.keySet()) {
            heap.offer(key);
            if (heap.size() > k) heap.poll();         // evict the weakest
        }
        List<Integer> res = new ArrayList<>(heap);    // heap order, NOT sorted
        // if sorted output is needed: poll repeatedly and reverse
        return res;
    }

    // 2) Arrival + ready-queue simulation (Single-Threaded CPU style).
    // tasks[i] = {enqueueTime, processingTime}; returns the processing order of indices.
    static int[] schedule(int[][] tasks) {
        int n = tasks.length;
        Integer[] idx = new Integer[n];
        for (int i = 0; i < n; i++) idx[i] = i;
        Arrays.sort(idx, Comparator.comparingInt(i -> tasks[i][0]));

        PriorityQueue<Integer> ready = new PriorityQueue<>((a, b) ->
            tasks[a][1] != tasks[b][1] ? Integer.compare(tasks[a][1], tasks[b][1])
                                       : Integer.compare(a, b));   // tie-break on index
        int[] order = new int[n];
        long time = 0;
        int p = 0, out = 0;
        while (out < n) {
            if (ready.isEmpty() && time < tasks[idx[p]][0]) time = tasks[idx[p]][0]; // idle: jump
            while (p < n && tasks[idx[p]][0] <= time) ready.offer(idx[p++]);
            int cur = ready.poll();
            order[out++] = cur;
            time += tasks[cur][1];
        }
        return order;
    }
}

Complexity

PriorityQueue.remove(Object) is O(n); if you must delete arbitrary items, use lazy deletion (skip stale tops) or a TreeMap.

Variations

• Top-K by value or score - size-k heap. Top K Frequent Elements, K Closest Points to Origin (max-heap on squared distance).
• Kth in a stream - persistent size-k min-heap. Kth Largest Element in a Stream.
• Resource allocation greedy - keep the k largest costs for the strong resource. Furthest Building You Can Reach (ladders on the biggest climbs, bricks on the rest).
• Event-driven scheduling - sort by arrival, heap by priority, jump the clock when idle. Single-Threaded CPU, Process Tasks Using Servers (two heaps: free servers and busy servers by finish time).
• K-way merge - one frontier element per list. Merge k sorted lists, kth smallest in a sorted matrix.
• Two heaps - balance lower and upper halves. Find median from data stream, sliding window median.

Common mistakes

• Comparators written as (a, b) -> a - b overflow for large or negative values; use Integer.compare / Long.compare or Comparator.comparingInt.
• Expecting iteration or toArray() over a PriorityQueue to be sorted - only repeated poll() returns elements in order.
• Changing an object's priority field while it is inside the heap; the heap is not re-ordered. Remove and re-insert, or push a new entry and skip stale ones.
• Forgetting the idle-time jump in simulations, which causes an infinite loop or NullPointerException when the ready-queue is empty.
• Not breaking ties the way the problem demands (smaller index, smaller id) - the expected order silently differs.

All practice problems

6 problems on this site use this pattern, easiest first. Solve them in order and say the instinct out loud before you code.