You almost never need the data sorted. You need its extremes. A heap of size K gives you those for O(n log K) instead of O(n log n), and it works on a stream you cannot hold in memory.
Sorting to find the top three of a million items does a million times more ordering work than the question asked for. A heap is the data structure that answers “what is the smallest thing I am currently keeping” in O(1) and lets you replace it in O(log K). That single operation is the whole pattern.
n with a small k.| Signal | What it looks like |
|---|---|
| K language | “top K”, “K-th largest”, “K most frequent”, “K nearest”. |
| Streaming | “as numbers arrive”, “design a class with an add method”, “you cannot store everything”. |
| Repeated extremes | “merge K sorted lists”, “the next smallest each time”, “schedule the earliest-finishing task”. |
| Two-sided balance | “median of a data stream” needs two heaps facing each other. Worth knowing as the classic follow-up. |
| Approach | Time | Space | Notes |
|---|---|---|---|
| Sort everything, take K | O(n log n) | O(n) | Simplest. Fine when n is small or K is near n. |
| Min-heap of size K | O(n log K) | O(K) | Works on a stream. The default answer. |
| Heapify all, pop K times | O(n + K log n) | O(n) | Better when K is tiny and you already hold all the data. |
| Quickselect | O(n) average | O(1) | O(n²) worst case. Needs mutable, in-memory data. |
| Bucket by count | O(n) | O(n) | Only when the key is a bounded integer, such as a frequency. |
heapqPython’s heapq is a min-heap over a plain list. Four facts cover almost everything.
| Need | Do this |
|---|---|
| A max-heap | Push -value, negate again on the way out. There is no max-heap in the standard library. |
| Build from a list | heapq.heapify(items) is O(n), cheaper than n pushes at O(n log n). |
| Pop then push, as one action | heapq.heapreplace(heap, item). One sift instead of two. |
| Order by a key | Push tuples: (sort_key, tiebreak, payload). Tuples compare left to right. |
TypeError. Always make the second element something comparable, such as an insertion counter or a coordinate, and keep unorderable payloads in third place or later.import heapq
# Each entry is (priority, counter, task_name). Tuples compare left to right.
heap: list[tuple[int, int, str]] = []
# priority sorts first. counter is unique, so a tie never reaches task_name.
heapq.heappush(heap, (priority, counter, task_name)) # counter breaks ties safely
heapq.nlargest(k, iterable) and nsmallest do exactly this pattern internally, and are the right production choice. In an interview, write the heap loop, then say “in real code I would call heapq.nlargest”. Both halves of that sentence score.import heapq
def k_largest(nums: list[int], k: int) -> list[int]:
"""The K largest values, unordered. Min-heap: the root is the weakest kept."""
# Starts empty and never holds more than k values.
heap: list[int] = []
# Invariant: heap holds the k largest values seen so far (fewer at the start).
for value in nums:
if len(heap) < k:
# Still filling up: every value earns a seat until there are k.
heapq.heappush(heap, value)
# heap[0]: index 0 is the min-heap root, the smallest value kept.
elif value > heap[0]: # beats the weakest survivor
# Pop the root and push value in one sift. Size stays at k.
heapq.heapreplace(heap, value)
# Only heap[0] is ordered. The rest of the list is in heap order, not sorted.
return heap
For the K smallest, flip both: use a max-heap by negating, and admit when value < -heap[0].
Reading the figure. Amber is the incoming value. Blue boxes change the heap. Red means the value is dropped. Notice that only one test, value > heap[0], runs once the heap is full.
def k_best(items: list[str], k: int) -> list[str]:
"""K items with the largest score(item)."""
# Entries are (score, index, item). Min-heap on score: root is the weakest kept.
heap: list[tuple[float, int, str]] = []
# Invariant: heap holds the k best-scoring items seen so far.
for index, item in enumerate(items):
# score(item): placeholder for the ranking key, bigger is better.
# Top K Frequent: score is the count. K Closest: score is -(x*x + y*y).
entry = (score(item), index, item) # index is the safe tie-break
if len(heap) < k:
# Fewer than k kept: admit without a contest.
heapq.heappush(heap, entry)
# heap[0]: index 0 is the root, the weakest entry kept.
elif entry > heap[0]:
# Swap out the weakest for the new entry. Size stays at k.
heapq.heapreplace(heap, entry)
# Unpack each tuple. The two _ skip score and index, keeping only item.
return [item for _, _, item in heap]
Reading the figure. Python compares tuples one slot at a time, left to right. Amber shows the tie on score. Green shows the index that settles it. The grey item is never looked at. The red bar shows what goes wrong with no index.
heapreplace. Correct but twice the sifting, and it briefly holds K+1 items.heap[0] is meaningful. Printing the list to “check” it will mislead you.Return the k-th largest element in an unsorted array. Note this is the k-th largest in sorted order, not the k-th distinct value.
k largest values seen. When the scan ends, the root is the smallest of the top k, which is precisely the k-th largest.k values and heapify them in O(k), rather than pushing them one at a time.sorted(nums)[-k] as the one-liner.import heapq
def find_kth_largest(nums: list[int], k: int) -> int:
"""The k-th largest value in nums, counting duplicates.
Args:
nums: Integers, unsorted.
k: 1-based rank from the top, with 1 <= k <= len(nums).
Returns:
The k-th largest value.
Raises:
ValueError: If k is out of range.
Example:
>>> find_kth_largest([3, 2, 1, 5, 6, 4], 2)
5
"""
# 1: k is a 1-based rank, so the smallest legal k is 1. len(nums) is the largest.
if not 1 <= k <= len(nums):
raise ValueError(f"k must be in 1..{len(nums)}, got {k}")
# Seed with the first k values. heapify is O(k), k pushes would be O(k log k).
heap = nums[:k] # slice 0..k-1: exactly k values
heapq.heapify(heap)
# Invariant: heap holds the k largest of the values scanned so far.
for value in nums[k:]: # k onward: the values not yet seeded
# heap[0]: index 0 is the root, the smallest of the k kept.
if value > heap[0]:
# Evict the root, insert value. Size stays at k.
heapq.heapreplace(heap, value)
# The root is the weakest of the k largest, which is the k-th largest.
return heap[0] # index 0: the min-heap root
nums = [3, 2, 1, 5, 6, 4], k = 2. Seed the heap with the first two values. Then scan the other four.
Reading the figure. Each frame is one step of the loop. The top circle is the root, the weakest value kept. Amber is a value that gets in. Red is a value that is skipped or evicted. The green root in the last frame is the answer.
import random
def _partition(nums: list[int], left: int, right: int, pivot_index: int) -> int:
"""Lomuto partition. Returns the pivot's final resting index."""
pivot = nums[pivot_index]
# Park the pivot at the right end so the scan below never moves it.
nums[pivot_index], nums[right] = nums[right], nums[pivot_index]
# store is the next slot for a value smaller than the pivot.
store = left
# range stops before right: right holds the parked pivot.
# Invariant: nums[left..store-1] are all less than pivot.
for i in range(left, right):
if nums[i] < pivot:
# Move the small value into the "less than" zone.
nums[store], nums[i] = nums[i], nums[store]
store += 1 # +1: the zone grew by one slot
# Drop the pivot between the two zones. It is now in its sorted position.
nums[store], nums[right] = nums[right], nums[store]
return store
def find_kth_largest_quickselect(nums: list[int], k: int) -> int:
"""O(n) expected time. Mutates nums.
The k-th largest sits at index len(nums) - k once the array is
partitioned around it, so we only recurse into the half that contains
that index.
"""
# In ascending order the k-th largest sits k slots from the end.
target = len(nums) - k
# 0 is the first index. len(nums) - 1 is the last index.
left, right = 0, len(nums) - 1
# Invariant: target lies inside nums[left..right], the unsettled window.
while left <= right:
# A random pivot is what turns the O(n^2) worst case into a
# vanishingly unlikely one.
pivot_index = _partition(nums, left, right, random.randint(left, right))
if pivot_index == target:
# The pivot landed exactly on the rank we want.
return nums[pivot_index]
if pivot_index < target:
# Target is to the right. + 1 skips the settled pivot.
left = pivot_index + 1
else:
# Target is to the left. - 1 skips the settled pivot.
right = pivot_index - 1
# Only reached when k was outside 1..len(nums).
raise ValueError("k out of range")
T(n) = 2T(n/2) + n, which is O(n log n). Quickselect recurses into one half: T(n) = T(n/2) + n, and the geometric series n + n/2 + n/4 + … sums to 2n. That one-sentence comparison is the answer the interviewer is fishing for.| Situation | Pick |
|---|---|
| Data arrives as a stream, or n is huge and k is small | Heap |
| All data in memory, mutation allowed, k arbitrary | Quickselect |
| You need a guaranteed worst case | Heap. Quickselect degrades to O(n²). |
| Production Python | heapq.nlargest(k, nums)[-1] |
k == 1: the maximum. k == len(nums): the minimum.[3, 3, 3] with k = 2 gives 3. Rank counts positions, not distinct values.Given an array nums and an integer k, return the k most frequent elements, in any order.
Counter in O(n). The selection is where the choice lives.k highest counts, O(m log k) where m is the number of distinct values.n, so index an array of lists by count and walk it downward. That is O(n) and beats the heap. Interviewers who ask for “better than O(n log k)” are asking for this.import heapq
from collections import Counter
def top_k_frequent(nums: list[int], k: int) -> list[int]:
"""The k most frequent values in nums, in no particular order.
Args:
nums: Input values.
k: How many to return, with 1 <= k <= number of distinct values.
Returns:
The k values with the highest counts.
Example:
>>> sorted(top_k_frequent([1, 1, 1, 2, 2, 3], 2))
[1, 2]
"""
# Phase 1: count each value in one O(n) pass.
counts = Counter(nums)
# Frequency first, so the min-heap root is the least frequent value kept.
heap: list[tuple[int, int]] = [] # (frequency, value)
# Invariant: heap holds the k highest-count values seen so far.
for value, frequency in counts.items():
if len(heap) < k:
# Fewer than k kept: admit without a contest.
heapq.heappush(heap, (frequency, value))
# heap[0][0]: [0] is the root entry, the second [0] is its frequency.
elif frequency > heap[0][0]:
# More frequent than the weakest kept: swap it in. Size stays at k.
heapq.heapreplace(heap, (frequency, value))
# The _ drops the frequency and keeps only the value.
return [value for _, value in heap]
from collections import Counter
def top_k_frequent_buckets(nums: list[int], k: int) -> list[int]:
"""Same answer in linear time, by indexing buckets on the count itself.
A value cannot appear more than len(nums) times, so counts are bounded
and can be used directly as array indices. That removes the log factor.
"""
counts = Counter(nums)
# buckets[f] holds every value that occurs exactly f times.
# len(nums) + 1 slots so index len(nums) exists: counts run 1..len(nums).
# A comprehension gives each slot its own list. [[]] * n would share one.
buckets: list[list[int]] = [[] for _ in range(len(nums) + 1)]
for value, frequency in counts.items():
buckets[frequency].append(value) # the count itself is the index
out: list[int] = []
# Walk counts from highest to lowest. Start at len(nums), the largest count.
# Stop value 0 is excluded, so the last count tried is 1. -1 steps downward.
# Invariant: out holds values with counts above the current frequency.
for frequency in range(len(nums), 0, -1):
for value in buckets[frequency]:
out.append(value)
if len(out) == k: # k collected: the most frequent k
return out
# Reached only when k exceeds the number of distinct values.
return out
nums = [1, 1, 1, 2, 2, 3], k = 2. Counts are {1: 3, 2: 2, 3: 1}. Buckets: index 1 holds [3], index 2 holds [2], index 3 holds [1]. Walking down from index 6: nothing at 6, 5, 4; at 3 take 1; at 2 take 2; we now have two, so return [1, 2].
Reading the figure. Blue cells hold counts or values not yet taken. Grey cells are empty buckets. The amber arrow is the walk from high counts to low. Green buckets are the ones the walk takes before it reaches k.
def top_k_frequent_stdlib(nums: list[int], k: int) -> list[int]:
"""What you would actually ship."""
# most_common(k) returns (value, count) pairs, highest count first.
# The _ drops the count and keeps only the value.
return [value for value, _ in Counter(nums).most_common(k)]
most_common uses heapq.nlargest internally when k is given, so it is the heap solution with none of the code. Write your own first, then show this.
k equals the number of distinct values: returns all of them.k encountered are returned.k = 0: returns []. Ask whether k can be zero.Given a list of points on a plane and an integer k, return the k points closest to the origin. The distance is the usual Euclidean distance.
sqrt is monotone increasing, so ordering by x² + y² gives exactly the same ranking as ordering by √(x² + y²). Skipping it removes n floating-point operations and, more importantly, keeps the comparisons in exact integer arithmetic. Say this unprompted.import heapq
def k_closest(points: list[list[int]], k: int) -> list[list[int]]:
"""The k points nearest the origin, in no particular order.
Args:
points: [x, y] integer coordinates.
k: How many to return, with 1 <= k <= len(points).
Returns:
The k closest points.
Example:
>>> sorted(k_closest([[1, 3], [-2, 2], [5, 8], [0, 1]], 2))
[[-2, 2], [0, 1]]
"""
# Negated squared distance turns heapq's min-heap into a max-heap, so the
# root is the furthest point we are currently keeping. x and y are the
# tie-break, which keeps every comparison on plain integers.
heap: list[tuple[int, int, int]] = []
# Invariant: heap holds the k closest points seen so far.
for x, y in points:
# x*x + y*y: squared distance, same order as the true distance.
# The leading minus flips the order: the furthest point becomes the root.
entry = (-(x * x + y * y), x, y)
if len(heap) < k:
# Fewer than k kept: admit without a contest.
heapq.heappush(heap, entry)
# entry[0]: its negated distance. heap[0][0]: root entry, its negated distance.
# Bigger negated distance means smaller real distance.
elif entry[0] > heap[0][0]: # closer than the worst one kept
# Evict the furthest kept point. Size stays at k.
heapq.heapreplace(heap, entry)
# The _ drops the negated distance. x and y rebuild the point.
return [[x, y] for _, x, y in heap]
points = [[1,3],[-2,2],[5,8],[0,1]], k = 2. Squared distances are 10, 8, 89, 1.
| point | d² | heap root (furthest kept) | action |
|---|---|---|---|
| [1, 3] | 10 | — | heap not full, push |
| [-2, 2] | 8 | — | heap not full, push |
| [5, 8] | 89 | 10 | 89 > 10, reject |
| [0, 1] | 1 | 10 | 1 < 10, evict [1, 3] |
Result: [[-2, 2], [0, 1]].
Reading the figure. Each frame handles one point. Boxes inside a frame are the heap, with the root on top and a thick border. Amber is a point that gets in, red is a point that is skipped or evicted. Green marks the final two points.
| Approach | Time | When |
|---|---|---|
| Sort by squared distance, slice k | O(n log n) | Simplest. Fine unless n is very large. |
| Max-heap of size k | O(n log k) | Streaming points, or n far bigger than k. |
| Quickselect on squared distance | O(n) average | All points in memory and you want the best average time. |
heapq.nsmallest(k, points, key=...) | O(n log k) | Production. |
k == len(points): every point is returned.(d, x, y) tuple makes the comparison total so nothing raises.0, always closest.heapq. Negated priorities work on older Python versions. Use heapify for linear-time heap construction.sqrt. They cost time and change no ordering.Pattern 10, Subsets and Backtracking, goes the other way. Instead of shrinking the problem to K items, it enumerates an exponential space on purpose, and the skill is pruning it before it eats the clock.