templates · Heaps & tries

Top k with a size-k heap

  • "The k largest / smallest / most frequent / closest"
  • "The k-th largest" in an array or a stream of numbers
  • k is much smaller than n, or the data arrives one item at a time

Top-k with a size-k heap

the one rule

The invariant

The heap always holds the k best items seen so far, and its top is the weakest of them.

write it from memory

The template

k-th largest

javajava.util.*
/** k-th largest: keep the k largest seen so far in a MIN-heap; its top is the smallest of them. */static int kthLargest(int[] nums, int k) {    PriorityQueue<Integer> heap = new PriorityQueue<>();   // min-heap of the k best    for (int x : nums) {        heap.offer(x);        if (heap.size() > k) heap.poll();                  // too many: drop the weakest    }    return heap.peek();}

For the k largest, use a min-heap: the top is the one to throw out when a better item arrives. Opposite, but simple hai.

fill in the blanks — only these lines change per problem

  1. line 3Heap direction. Min-heap for "k largest"; max-heap (Collections.reverseOrder()) for "k smallest" or "k closest".
  2. line 6When to drop: as soon as the size goes over k.
  3. line 8What to return: the top (k-th best), or everything in the heap (the top k).

same skeleton: top k frequent

javajava.util.*
/** The same skeleton on (value, count) pairs: the k most frequent values, any order. */static int[] topKFrequent(int[] nums, int k) {    Map<Integer, Integer> count = new HashMap<>();    for (int x : nums) count.merge(x, 1, Integer::sum);    PriorityQueue<Map.Entry<Integer, Integer>> heap =        new PriorityQueue<>((a, b) -> Integer.compare(a.getValue(), b.getValue())); // weakest = fewest    for (Map.Entry<Integer, Integer> e : count.entrySet()) {        heap.offer(e);        if (heap.size() > k) heap.poll();    }    int[] out = new int[heap.size()];    for (int i = out.length - 1; i >= 0; i--) out[i] = heap.poll().getKey(); // most frequent first    return out;}

Only the comparator changed. Count first with merge(x, 1, Integer::sum), then run the size-k heap over the entries.

fill in the blanks — only these lines change per problem

  1. line 6What "weakest" means: fewest count, farthest distance, or a tie-break on the word for "top k frequent words".

complexity

What it costs

time O(n log k)
Each of the n items costs one offer and at most one poll on a heap of size k + 1, and each of those is O(log k). The frequency version first counts in O(n), then runs the heap over d distinct values: O(n + d log k).
space O(k)
The heap never holds more than k + 1 items. The frequency version also needs O(d) for the HashMap of d distinct values.

watch out

Common bugs

  • A comparator written as a - b. It overflows for large or negative values; use Integer.compare(a, b).
  • Using a max-heap of all n items for "k largest". It works, but it is O(n log n) time and O(n) space.
  • Expecting PriorityQueue to iterate in sorted order. Only poll() gives sorted order.

use it

Practise it on

learn it properly