Top K Elements with Heaps: A Practical Interview Pattern

Why Top K Shows Up So Often

If you do enough coding interviews, you will keep seeing variations of the same idea: find the top K elements. Sometimes it is the kth largest element, sometimes k closest points to the origin, sometimes the top K frequent items. The surface details change, but the core decision is the same: how do you keep only the most relevant K candidates while scanning a much larger input?

This pattern matters because it tests more than syntax. Interviewers want to see whether you can recognize when a full sort is wasteful, when a heap gives you the right tradeoff, and when a specialized approach like quickselect is better. If you learn the pattern once, you can reuse it across a lot of problems.

Problem Statement

The general top K problem looks like this:

• You are given a collection of items.
• You need to return the K best items according to some ranking rule.
• “Best” could mean largest, smallest, closest, most frequent, or highest score.

Examples:

• Return the kth largest number in an array.
• Return the K closest points to the origin.
• Return the top K frequent words.

The key challenge is that the input may be large, so you should avoid sorting everything unless that is truly necessary.

Intuition: Why Heaps Fit This Pattern

A heap is useful when you repeatedly need access to the smallest or largest item.

For top K problems, the trick is usually to keep a heap of size at most K:

• If you want the K largest items, keep a min-heap of size K.
• If you want the K smallest items, keep a max-heap of size K.

Why this works:

• The heap stores only the current best K candidates.
• When a new item arrives, compare it against the “worst” item in your heap.
• If the new item is better, replace the worst one.

This gives you a streaming-style solution: you can process the input one item at a time without sorting the whole thing.

Mental model

Think of the heap as a small gatekeeper:

• It lets in only the best K candidates.
• It pushes out the least useful candidate whenever a better one arrives.

That simple rule is the core of most top K heap solutions.

Approach 1: Top K Largest with a Min-Heap

Suppose you want the K largest numbers in an array. Use a min-heap:

• Push each number into the heap.
• If the heap grows beyond size K, pop the smallest one.
• At the end, the heap contains the K largest numbers.
• The smallest element in the heap is the kth largest overall.

Example 1: kth largest element

import heapq

def kth_largest(nums, k):
    heap = []

    for num in nums:
        heapq.heappush(heap, num)
        if len(heap) > k:
            heapq.heappop(heap)

    return heap[0]

print(kth_largest([3, 2, 1, 5, 6, 4], 2))  # 5

Walkthrough

For nums = [3, 2, 1, 5, 6, 4], k = 2:

1. Push 3 → heap = [3]
2. Push 2 → heap = [2, 3]
3. Push 1 → heap = [1, 3, 2], size is now 3, so pop smallest → heap = [2, 3]
4. Push 5 → heap = [2, 3, 5], pop smallest → heap = [3, 5]
5. Push 6 → heap = [3, 5, 6], pop smallest → heap = [5, 6]
6. Push 4 → heap = [4, 6, 5], pop smallest → heap = [5, 6]

The heap now contains the 2 largest numbers: [5, 6]. The smallest among them, 5, is the 2nd largest.

Why this is good

You never sort the full array. You only maintain a heap of size K.

Approach 2: Top K Smallest with a Max-Heap

Python’s heapq module only provides a min-heap directly. To simulate a max-heap, store negative values.

If you want the K smallest numbers:

• Push -num into the heap.
• Keep the heap size at most K.
• The number with the largest magnitude negative value corresponds to the smallest original number.

Example 2: K smallest numbers

import heapq

def k_smallest(nums, k):
    heap = []

    for num in nums:
        heapq.heappush(heap, -num)
        if len(heap) > k:
            heapq.heappop(heap)

    return [-x for x in heap]

print(k_smallest([7, 10, 4, 3, 20, 15], 3))  # order may vary: [7, 4, 3]

Walkthrough

For nums = [7, 10, 4, 3, 20, 15], k = 3:

• Store negatives: -7, -10, -4, -3, -20, -15
• The heap keeps the three most negative values, which correspond to the three smallest numbers.

This is less common in interviews than the min-heap version, but the idea is identical.

Approach 3: K Closest Points to the Origin

This is a classic top K problem, but the ranking rule is not based on the number itself. Instead, compare squared distance to the origin:

• Distance squared for point (x, y) is x*x + y*y
• You do not need the square root, because it does not change ordering

Use a max-heap of size K so you can discard the farthest point among the current best K.

Example 3: k closest points

import heapq

def k_closest(points, k):
    heap = []

    for x, y in points:
        dist = x * x + y * y
        heapq.heappush(heap, (-dist, x, y))
        if len(heap) > k:
            heapq.heappop(heap)

    return [[x, y] for _, x, y in heap]

points = [[1, 3], [-2, 2], [5, 8], [0, 1]]
print(k_closest(points, 2))  # [[-2, 2], [0, 1]] in some order

Walkthrough

For points = [[1, 3], [-2, 2], [5, 8], [0, 1]], k = 2:

• (1, 3) → distance 10
• (-2, 2) → distance 8
• (5, 8) → distance 89
• (0, 1) → distance 1

The heap stores (-distance, x, y) so the point with the largest distance is easiest to remove.

After processing all points, the heap contains the two closest points.

Step-by-Step Template You Can Reuse

When you see a top K problem, ask these questions:

1. What is the ranking rule?

   • Largest value?
   • Smallest value?
   • Closest distance?
   • Highest frequency?

2. Do I need the items sorted, or just the best K?

   • If you only need the best K, a heap is often enough.

3. Which heap size should I maintain?

   • Usually size K.

4. Should I use a min-heap or a max-heap?

   • K largest → min-heap
   • K smallest → max-heap
   • K closest → max-heap on distance

5. What should I store in the heap?

   • The value itself
   • A tuple like (priority, value)
   • A transformed value like -num

General Python Template

Here is a reusable pattern for many heap-based top K problems:

import heapq

def top_k(items, k, score_fn):
    heap = []

    for item in items:
        score = score_fn(item)
        heapq.heappush(heap, (score, item))
        if len(heap) > k:
            heapq.heappop(heap)

    return [item for _, item in heap]

This version assumes smaller scores are worse, so it keeps the K largest by score. If you want the opposite, invert the score or use negative values.

For example:

def k_highest_scores(records, k):
    return top_k(records, k, lambda r: r[0])

That template is not always the final answer, but it is a good starting point when the ranking is custom.

Complexity Analysis

Different approaches solve top K problems with different tradeoffs.

Heap vs quickselect

Heaps are usually easier to code correctly under interview pressure. Quickselect can be faster on average, but it is more delicate:

• It partitions the array around a pivot.
• It only recurses into the side containing the target.
• It is excellent for kth largest/smallest, but less natural for problems with custom ranking or streaming input.

If the question is “return the kth largest element,” quickselect is worth knowing. If the question is “return the top K items,” heaps are often the cleaner solution.

When to Prefer Heap Over Quickselect

Use a heap when:

• You need the top K items, not just the kth item.
• The input arrives as a stream.
• You need to process data incrementally.
• The ranking depends on a custom key.

Use quickselect when:

• You only need the kth element.
• You want average linear time.
• You are comfortable implementing partition logic.

Related Problem: Top K Frequent Elements

Frequency-based problems are also top K problems. You count first, then rank by frequency.

import heapq
from collections import Counter

def top_k_frequent(nums, k):
    freq = Counter(nums)
    heap = []

    for num, count in freq.items():
        heapq.heappush(heap, (count, num))
        if len(heap) > k:
            heapq.heappop(heap)

    return [num for _, num in heap]

The pattern is the same:

• Build a score for each candidate.
• Keep only the best K scores in a heap.

Common Mistakes

A few mistakes show up repeatedly:

• Using a full sort when a heap is enough

  • Sorting works, but it costs O(n log n).

• Forgetting to cap the heap at size K

  • If you never pop, the heap does not solve the top K problem.

• Using the wrong heap direction

  • K largest → min-heap of size K
  • K smallest → max-heap of size K

• Using sqrt for distances

  • For comparison, squared distance is sufficient.

• Returning heap order as final sorted order

  • A heap is not fully sorted. If order matters, you may need to sort the final K items.

Key Takeaways

• Top K problems usually ask you to keep only the most relevant K candidates.
• A heap of size K is the standard tool.
• Use a min-heap for K largest items and a max-heap for K smallest items.
• For custom ranking, store (score, item) or a transformed score.
• Quickselect is a strong alternative for kth-element problems, but heaps are more flexible.

Practice These Next

Try these problems to reinforce the pattern:

1. Kth Largest Element in an Array
2. K Closest Points to the Origin
3. Top K Frequent Elements
4. Find Median from Data Stream
5. Kth Smallest Element in a Sorted Matrix

If you'd like to see the variables update live as you step through, check out codebrew.ai.