Top-k elements — LeetCode patterns

Problem

You need the k best items by some measure: the k largest numbers, the k most frequent words, the k points closest to the origin, or just the k-th largest. The data may arrive as a stream you can see only once, or be too large to sort comfortably. In all cases the answer has size k, and k is usually far smaller than n. What you want is a structure of size k, not size n.

Signal	Pattern	Why
"The k largest" or "k-th largest"	Min-heap of size k	The root is the k-th largest so far and the bar a newcomer must beat.
"The k smallest" or "k closest"	Max-heap of size k	The root is the worst of the kept items, the one to evict first.
"The k most frequent"	Count with a dictionary, then a heap of size k (or buckets)	Counting gives one score per distinct item, and the heap keeps the k best.
"Stream: keep the k-th largest after every number"	Heap that survives between calls	Each arrival is one O(log k) update, with no re-sort.
"k is close to n"	Sort instead	A heap of size nearly n does about n log n work and sorting is simpler.
"Only the k-th value is needed, input is static"	Quickselect	Average O(n) beats O(n log k) when the data can be rearranged in place; a later page covers it.
"Scores are small whole numbers"	Bucket by score	Counting sort gives O(n) without any heap, as the counting page showed.

Intuition

Imagine a podium with k places and a doorkeeper who always knows who stands on the lowest step. The first k visitors step straight up. After that, a visitor only gets in by beating the person on the lowest step, who is then sent away. Visitors who do not beat the lowest step are turned away at the door at once. At the end the podium holds exactly the k best, because anyone who was ever evicted was beaten by k other people who stayed or were themselves beaten by better ones.

The surprise is the direction of the heap. To keep the largest k you use a min-heap: the item you need quick access to is the weakest of the strong, since that is the one to compare against and to throw out. Keeping the smallest k uses a max-heap for the same reason. Getting this backwards, a max-heap for the largest, is the classic mistake, and it forces a full sort.

Green bars are the 3 largest. The dashed line marks the weakest kept value: a newcomer must exceed it to get in. The diagram scrolls sideways on a narrow screen.

Mechanics

The invariant

After processing any prefix of the input, the heap holds exactly the k best items of that prefix (or all of them, if fewer than k have been seen), and its root is the weakest of those. This is true at the start (empty) and survives each step. If the heap has room, adding the item keeps it true. If the heap is full and the newcomer beats the root, then the root is not among the k best of the longer prefix, since k items now beat it, so replacing it keeps the invariant. If the newcomer does not beat the root, then at least k kept items are at least as good, so the newcomer is not in the k best and is dropped. When the input ends, the heap is the answer, and its root is the k-th best.

Why O(n log k)

Each of the n items costs one comparison with the root, and at most one push or replace, each O(log k) because the heap never grows beyond k. Sorting costs O(n log n). When k is small the gap is large: for a million items and k = 10, a heap does about log₂ 10 ≈ 3 levels of work per item, a sort about 20. The heap also needs only O(k) memory, so it works on a stream that does not fit in memory.

Frequency and distance variants

The heap only needs a score for each item. For frequency, first count with a dictionary (one pass, O(n)), then run the same procedure over the distinct items with the count as the score. For distance, the score is the distance to a point; compare squared distances, which rank the same and avoid a square root. To keep the closest, you want to evict the farthest, so the heap root is the worst kept item and the heap is a max-heap on distance, usually stored as negated values in Python.

The bar is the weakest value kept. It rises only when a replacement happens, which is why the scan gets cheaper as it goes: more and more values are skipped with a single comparison.

Worked example

Find the 3 largest of 4, 9, 1, 7, 3, 8, 2, 6. Each row shows the action, the bar (the heap root) and the values kept.

Value	Action	Bar after	Kept, largest first
4	fill	4	4
9	fill	4	9, 4
1	fill	1	9, 4, 1
7	replace 1	4	9, 7, 4
3	skip	4	9, 7, 4
8	replace 4	7	9, 8, 7
2	skip	7	9, 8, 7
6	skip	7	9, 8, 7

The first 3 values fill the heap. Then 7 beats the bar 1 and evicts it; 3 does not beat the bar 4 and is skipped. The answer is 9, 8, 7, and the 3rd largest is the final bar, 7.

Most frequent words. For the day is sunny the the the sunny is is the counts are the 4, day 1, is 3, sunny 2, from 4 distinct words among 10. The top 2 are the and is.

Closest points. To keep the 2 points closest to the origin, compare squared distances:

Point	Squared distance	Action	Worst kept
(3, 3)	18	fill	18
(5, -1)	26	fill	26
(-2, 4)	20	replace	20
(1, 1)	2	replace	18
(0, 2)	4	replace	4

The answer is (0, 2) and (1, 1). The worst kept distance only falls once the heap is full, which is the same shape as the bar above, mirrored.

Implementation

import heapq
from collections import Counter


def largest_k(nums, k):
    heap = []
    for x in nums:
        if len(heap) < k:
            heapq.heappush(heap, x)
        elif x > heap[0]:
            heapq.heapreplace(heap, x)      # pop the weakest and push x in one step
    return sorted(heap, reverse=True)


class KthLargest:
    def __init__(self, k, nums):
        self.k = k
        self.heap = []
        for x in nums:
            self.add(x)

    def add(self, x):
        if len(self.heap) < self.k:
            heapq.heappush(self.heap, x)
        elif x > self.heap[0]:
            heapq.heapreplace(self.heap, x)
        return self.heap[0]


def top_k_frequent_words(words, k):
    counts = Counter(words)
    # most frequent first, ties alphabetical: sort key (-count, word)
    return heapq.nsmallest(k, counts, key=lambda w: (-counts[w], w))


def k_closest(points, k):
    heap = []                                 # max-heap by distance, stored negated
    for x, y in points:
        d = x * x + y * y
        if len(heap) < k:
            heapq.heappush(heap, (-d, x, y))
        elif d < -heap[0][0]:
            heapq.heapreplace(heap, (-d, x, y))
    return [[x, y] for _, x, y in heap]

heapq.nlargest and heapq.nsmallest implement the same idea, and they take a key. Writing the loop by hand is worth it in an interview because the interviewer usually asks what happens inside.

Variations and edge cases

Ties. The word problem needs a deterministic order. Make the score a tuple, with the count negated so that smaller means better, and the word as the tiebreak. With the heap-of-size-k approach, a max-heap on (count, reversed word) is harder to write than nsmallest with a key.
Bucket by count. If scores are counts, an array of lists indexed by count, read from the high end, gives the top k in O(n) with no heap.
The k-th element only. A heap of size k gives the answer as its root; quickselect gives it in average O(n) time but changes the input.
Streams. Keep the heap as state between calls; each update is O(log k).
Edge cases: k equal to 0 (return empty, and guard heap[0]), k larger than n (return everything), duplicates in the input (equal values count separately for the largest k, so check what the problem means), negative numbers, and an empty heap when a stream is queried before k values arrive.

Complexity and tradeoffs

Time: O(n log k): n scans, each with at most one heap update on k items. Counting first, as for frequency, adds an O(n) pass, and then O(d log k) over d distinct items.
Space: O(k) for the heap, plus O(d) for a frequency table.
Baseline. Sort everything: O(n log n) time, O(n) space, three lines of code. It is the right answer when k is a large fraction of n or when n is small.
When it fails. Quickselect has better average time, O(n), for a one-off query on data you may rearrange. If you need the top k repeatedly under insertions and deletions, a balanced tree or a sorted container works better than a heap, which cannot delete arbitrary items cheaply. For small integer scores, buckets beat both.

Problems

The three home problems are the three flavours: a stream's k-th largest, the most frequent words with a tie rule, and the closest points by distance. The classic array version and the most frequent numbers appear below in "also uses", taught on neighbouring pages.

Easy

703. Kth Largest Element in a Stream NeetCode 150 NeetCode 250
A min-heap holding the k largest values has the kth largest on top, and each add only touches the top.

Medium

692. Top K Frequent Words Grind 169
Count words, then select by frequency descending with alphabetical order breaking ties, using a heap or a sort on that composite key.
973. K Closest Points to Origin Grind 169 NeetCode 150 NeetCode 250
Keep a max-heap of the k closest points by squared distance, or quickselect on distance.

Also uses this pattern

These problems have a different home page, but this idea is part of their solution.

215. Kth Largest Element in an Array Medium
347. Top K Frequent Elements Medium

Connections

Read first

Heaps intro
Counting and bucket sort core

Check yourself

Why do we use a min-heap to keep the k largest values?

Each newcomer must be compared with the weakest kept value, and that is the item to evict if the newcomer wins. A min-heap exposes its smallest value at the root in O(1). In the example the root after the scan is 7, the 3rd largest.

Why can a skipped value never belong to the answer?

When the heap is full and a value does not beat the root, k kept values are at least as large, so the value ranks below k others and cannot be among the k best. Values such as 3 and 2 in the example were skipped for this reason.

When is sorting the better answer?

When k is close to n, the heap does about n log n work with a larger constant than the built-in sort, so a one-line sort is as fast and simpler. It is also right when n is tiny or the whole ordering is needed anyway.

Why compare squared distances for the closest points?

Squaring is increasing for non-negative numbers, so it ranks distances identically and avoids a square root, which also keeps every comparison in exact integer arithmetic. In the example (1, 1) has squared distance 2 and (0, 2) has 4.

References

Cormen, Leiserson, Rivest and Stein, Introduction to Algorithms, chapter 6 (priority queues) and chapter 9 (order statistics).
Sedgewick and Wayne, Algorithms, section 2.4: the "TopM" client of a priority queue.
Skiena, The Algorithm Design Manual, the section on priority queues and selection.
Python documentation, heapq: nlargest, nsmallest and the priority-queue recipes.
The NeetCode roadmap and the Blind 75, Grind 169, NeetCode 150 and NeetCode 250 lists, used as the index of topics and problems.