Top K Elements

Find the k largest, smallest, closest or most frequent items with a size-k heap, quickselect or frequency buckets, and know when each one is the right answer.

Overview

Top K problems ask for the k items that rank highest by some measure: the largest values, the most frequent words, the closest points. Sorting everything costs O(n log n) and throws most of the work away. Three tools do better: a heap of size k, quickselect, and frequency buckets.

The question behind a size-k heap

Keep the k best items seen so far in a heap. For every new item ask one question: is it better than the worst item I am keeping? To answer in O(1), the worst kept item must be at the root. For the k largest values, the worst kept value is the smallest one, so the heap is a min-heap. That is why the heap direction looks backward.

k largest — min-heap of size k — the smallest kept value — it is larger than the root
k smallest — max-heap of size k — the largest kept value — it is smaller than the root
k closest to a point — max-heap by distance, size k — the farthest kept point — it is closer than the root
k most frequent — min-heap by count, size k — the least frequent kept item — its count is higher than the root's

Trace for the 3 largest of 4, 9, 1, 7, 3, 8, 2. The heap is shown sorted; the root is its first value.

4 — empty — fewer than 3: push — 4
9 — 4 — push — 4, 9
1 — 4, 9 — push — 1, 4, 9
7 — 1, 4, 9 — 7 > 1: replace the root — 4, 7, 9
3 — 4, 7, 9 — 3 is not > 4: skip — 4, 7, 9
8 — 4, 7, 9 — 8 > 4: replace the root — 7, 8, 9
2 — 7, 8, 9 — 2 is not > 7: skip — 7, 8, 9

Each of the n values does at most one O(log k) heap operation, so the total is O(n log k) with O(k) memory. The heap never sees the whole input at once, which is why the same code works on a stream.

Three ways to get the k largest

Size-k heap — O(n log k) — O(k) — stream or array, unchanged — the k items; sorted if you pop them out
Quickselect — O(n) average, O(n^2) worst — O(1) — array, reordered in place — the k items, not sorted
Heapify all, pop k — O(n + k log n) — O(n) — array — sorted
Full sort — O(n log n) — depends on the sort — array — sorted

Pick by the situation: a stream or a small k favors the heap; a large in-memory array with a one-time query favors quickselect; k close to n makes a full sort as good as anything.

Quickselect: partition, then recurse on one side

Quickselect is quicksort that only follows the side containing the answer. Pick a pivot, move every smaller value to its left and every larger value to its right, and see where the pivot lands. If it lands at the index you want, you are done. Otherwise repeat on the one side that contains that index.

The k-th largest of n values is at sorted index n minus k. Trace for the 2nd largest of 3, 2, 1, 5, 6, 4 (target index 4), using the last value of each range as the pivot.

1 — 0 to 5 — 4 — 3, 2, 1, 4, 6, 5 — 3 — 3 < 4: go right
2 — 4 to 5 — 5 — 3, 2, 1, 4, 5, 6 — 4 — equal: answer 5

Each round discards part of the array, so the average work is n + n/2 + n/4 + ... which is O(n). The worst case is O(n^2) when every pivot is the smallest or largest value; picking the pivot at random makes that very unlikely.

Most frequent items: count, then bucket

Top K Frequent Elements needs a count first: one pass with a hash map, O(n). Then either run the size-k heap over (count, item) pairs, O(n log k), or use buckets. A count can never exceed n, so an array of n + 1 buckets indexed by count holds every item; read the buckets from the highest count down until you have k items.

The table shows the buckets for 1, 1, 1, 2, 2, 3 with k = 2.

3 — 1
2 — 2
1 — 3

Reading from the top gives 1, then 2: done. The bucket version is O(n) time and O(n) space; the heap version is O(n log k) time and O(k) space beyond the map. If two items tie at the k-th place, say how you break the tie; the problem usually promises that the answer is unique.

K closest points

Rank points by distance from the origin. The square root is never needed: comparing x^2 + y^2 gives the same order as comparing the distances, and it stays in integers. For the k closest, the worst kept point is the farthest one, so the heap is a max-heap by squared distance.

(3, 3) — 18 — (3, 3) — fewer than 2: push
(5, -1) — 26 — (5, -1), (3, 3) — push
(-2, 4) — 20 — (-2, 4), (3, 3) — 20 < 26: replace the root
(1, 1) — 2 — (3, 3), (1, 1) — 2 < 20: replace the root

Quickselect on squared distances gives the same answer in O(n) average when all points are in memory.

The k-th largest in a stream

Kth Largest Element in a Stream keeps the same size-k min-heap alive between calls. After each add, the root is the k-th largest value seen so far. The trace uses k = 3, starting from 4, 5, 8, 2.

start — 4, 5, 8 — 4
add 3 — 4, 5, 8 — 4
add 5 — 5, 5, 8 — 5
add 10 — 5, 8, 10 — 5
add 9 — 8, 9, 10 — 8
add 4 — 8, 9, 10 — 8

Each add is O(log k) and memory stays at O(k) no matter how long the stream runs.

Library shortcuts, and what they hide

k largest, sorted — heapq.nlargest(k, items) — a PriorityQueue of size k, then drain it
k smallest, sorted — heapq.nsmallest(k, items) — PriorityQueue with a reversed comparator
k most frequent — Counter(items).most_common(k) — a HashMap of counts plus a PriorityQueue on the entries

The heapq documentation says nlargest and nsmallest perform best for small k; for large k, sorted() is more efficient, and for k = 1 use min() or max(). Counter.most_common orders items with equal counts in the order they were first seen. Use the library call if it is allowed, and be ready to write the size-k heap by hand.

What you will be able to do

Explain why the k largest need a min-heap and the k smallest a max-heap, and pick the direction for any ranking.
Write the size-k heap for values, points and frequency pairs, and state O(n log k) time and O(k) space.
Write quickselect, trace a partition, and say when its O(n) average beats the heap.
Solve the frequency version in O(n) with buckets.
Keep the k-th largest of a stream with O(k) memory.

When to use it

The k largest or k smallest values of an array or a stream
The k most frequent items, words or characters
The k points, strings or records closest to a target by some distance
The k-th largest value, once (quickselect) or continuously (a size-k heap)
A partial ranking when a full sort would be wasted work

More Coding & DSA topics

Start free on Cariara · Coding problems · System design studio