Top-K Elements with a Heap — O(n log k) Pattern in Java

Top-K pattern finds the K largest/smallest/most-frequent elements in O(n log k) time and O(k) space using a heap.

The core idea

• Maintain a size-K heap of winners so far
• For each element: push; if size exceeds K, pop the worst
• Final heap = top K

Complexity

• Time: O(n log k) — each push/pop is O(log k)
• Space: O(k)

The counter-intuitive trick — opposite heap

Top-K LARGEST → MIN-heap

• Smallest at root — that's what we pop when adding a larger element
• Heap root = worst of our current top-K

Top-K SMALLEST → MAX-heap

• Largest at root — pop when adding a smaller element

Top-K MOST FREQUENT → min-heap keyed by frequency

Top-K MOST RECENT → min-heap keyed by timestamp

Reasoning: heap root is what we discard when exceeding K — must be the 'worst by ranking'.

Java PriorityQueue

• Default: min-heap
• Max-heap: `new PriorityQueue<>(Comparator.reverseOrder())`
• Ops: offer, poll, peek
• No random access; contains/remove are O(n)

Examples

Kth largest element

• Min-heap of size K
• peek() after iteration = Kth largest

Top-K frequent elements

• Step 1: HashMap of counts
• Step 2: min-heap of size K keyed by count
• Total: O(n + D log k) where D = distinct elements

K closest points to origin

• Max-heap of size K keyed by squared distance (no sqrt — ordering preserved)
• Top-K smallest distances → max-heap

Top-K words with tie-break

• Min-heap keyed by frequency
• Tiebreaker: larger lex at root (so it's popped)
• Result reversed at the end
• Tiebreaker direction is opposite of final desired ordering

Alternative approaches

Sort (O(n log n))

• Simplest
• Wasteful for small k
• Use for small n or simplicity

QuickSelect (O(n) average, O(n²) worst)

• Partition-based, like quicksort
• In-place, destroys input
• Fastest offline
• Randomize pivot to avoid worst case
• Doesn't work for streaming

Bucket sort (O(n))

• Only for bounded frequency range
• buckets[freq] = list of elements
• Linear time, O(n) space

Heap

• O(n log k) — best when k << n
• Works for streaming / online (only option)
• Generic comparator
• Simple to code

Decision matrix

Small k, offline → heap

Large k ≈ n/2, pure speed → QuickSelect

Streaming / online → heap

Bounded frequency range → bucket sort

Simplicity → sort

Common pitfalls

1. Using max-heap for top-K largest

• Wrong! Use min-heap
• Mental check: what needs to be popped when exceeding K? That's the heap root

2. Sqrt on distances

• Unnecessary; squared distance preserves ordering
• Avoids floating-point precision issues

3. Wrong tiebreaker direction

• Always think: 'what should be popped when two elements tie?'
• That's the worst by final ranking

4. Using contains on PriorityQueue

• O(n) — slow
• Use HashSet alongside if lookup needed

5. Not hinting initial capacity

• new PriorityQueue<>(k) avoids resize

6. Missing edge cases

• k > n: return all
• Empty input: return empty
• Ties: comparator must handle explicitly

Streaming top-K

• Heap is the only bounded-space option
• Maintain size-K heap as data flows in
• O(log k) per element amortized
• Used in real-time analytics, top-trends-now systems

PriorityQueue vs TreeMap

• PriorityQueue: fast for top-K, no random access
• TreeMap: sorted keys, ceiling/floor in O(log n), more flexible
• For top-K specifically, PriorityQueue is simpler

Recommendation: The top-K heap pattern is one of the most useful interview techniques — O(n log k) time, O(k) space. Memorize the counter-intuitive rule: use the opposite heap type — min-heap for top-K largest, max-heap for top-K smallest. The mental check: the heap root is what gets popped when we exceed K, so it must be the worst by the final ranking. For frequency problems (top-K frequent elements/words), combine HashMap for counts + heap of size K. For distance problems (K closest points), use squared distance (no sqrt) and a max-heap because we want the K smallest distances. For tiebreaker problems (top-K words with lex tiebreak), set the comparator so the element to discard is at the root — this often means the tiebreaker is in the opposite direction of the final ranking. Know the alternatives: QuickSelect is O(n) average offline but doesn't stream and has O(n²) worst case; sorting is O(n log n) — wasteful for small k; bucket sort is O(n) but needs bounded frequency range. Use heap when k << n or when data streams; use QuickSelect for large-k offline speed; use sorting for simplicity. In Java, PriorityQueue<>() is min-heap by default — use Comparator.reverseOrder() for max-heap. Avoid contains/remove(obj) (O(n)). The mental model: top-K = size-K heap of winners, pop worst when overflow, time O(n log k). Once you recognize the pattern (find K best of something), the code writes itself.