Top-K ProblemsTop K

01Why it exists

Reach for it when you see:Top K largest or smallest, the Kth largest, the K most frequent, the K nearest to some point, K far smaller than n, data arriving as a stream.

02The core idea

The obvious way to find the top K is to sort everything and take the first K, at O(n log n). But you only need K of them, and the order of the other n − K does not matter in the slightest. A min-heap of size K remembers just the best K seen so far, and its top is the smallest of those K — the cut-off for getting onto the list. A new element below the cut-off can be ignored outright; one above it evicts the cut-off and takes its place. Each element costs at most one O(log K) operation, which is O(n log K) overall.

The direction is easy to get backwards: the top K largest needs a min-heap, and the top K smallest needs a max-heap. The way to remember it is that the top of the heap is whatever is most likely to be evicted next. Once that same heap has scanned everything, its top is the Kth largest, which is why "the Kth largest" and "the top K largest" are really one problem.

Two other routes exist. Quickselect borrows quicksort's partition to throw away half the data at a time, O(n) on average, but it needs the whole dataset in memory and does not suit a stream. Bucket sort: when the thing being compared is a number of occurrences, that count is at most n, so open n buckets, drop each element into the bucket for its count, and collect from the largest bucket downwards, in O(n). Which one to reach for depends on how large K is, whether the data is a stream, and whether the values being compared have a bounded range.

03The algorithm

1. 1Fix the direction: top K largest wants a min-heap, top K smallest wants a max-heap. When the comparison is on a count or a distance, compute that value first and store (value, element) tuples in the heap.
2. 2Scan the elements one at a time. While the heap holds fewer than K, simply push.
3. 3Once it is full, compare each new element with the top: no better than the top means skip it (O(1)), and better than the top means pop the top and push the new one instead (Python's heapreplace does both in one go).
4. 4When the scan ends, the heap is the answer. For sorted output, pop repeatedly and reverse; if you only want the Kth largest, just read the top.
5. 5If K is close to n, the data is not a stream and you only need this once, quickselect or a plain sort may well be faster; when the comparison is on counts, consider bucket sort.

04Interactive demo

K = 3. The view counts are read one at a time, and the top of the heap (in amber) is the cut-off for getting onto the list. Notice how many values never touch the heap at all and are skipped outright.

05Code

Three variations — top K largest, the Kth largest, and the K most frequent — plus the O(n) bucket sort version for comparison. C++ also demonstrates nth_element, which is the standard library's quickselect.

import heapq
from collections import Counter


def top_k_largest(nums, k):
    """Top K largest: keep a min-heap of size K. O(n log k)"""
    heap = []
    for x in nums:
        if len(heap) < k:
            heapq.heappush(heap, x)
        elif x > heap[0]:                # only beating the threshold earns a place
            heapq.heapreplace(heap, x)   # pop the smallest, then push: one O(log k) step
    return sorted(heap, reverse=True)


def kth_largest(nums, k):
    """Kth largest: the same heap, and at the end the top of it is the answer"""
    heap = []
    for x in nums:
        heapq.heappush(heap, x)
        if len(heap) > k:
            heapq.heappop(heap)
    return heap[0]


def top_k_frequent(words, k):
    """Top K by frequency: count first, then run Top-K over (count, word)"""
    count = Counter(words)
    # nlargest is a size-k heap under the hood
    return heapq.nlargest(k, count, key=count.get)


def top_k_frequent_bucket(nums, k):
    """Counts never exceed n, so bucket sort brings this down to O(n)"""
    count = Counter(nums)
    buckets = [[] for _ in range(len(nums) + 1)]
    for x, c in count.items():
        buckets[c].append(x)
    out = []
    for c in range(len(buckets) - 1, 0, -1):
        out.extend(buckets[c])
        if len(out) >= k:
            return out[:k]
    return out

06Practice

• LeetCode 215Kth Largest Element in an Array (do it once with a heap and once with quickselect)Medium
• LeetCode 347Top K Frequent Elements (try the bucket sort version)Medium
• LeetCode 973K Closest Points to OriginMedium
• LeetCode 692Top K Frequent Words (ties broken lexicographically)Medium
• LeetCode 1985Find the Kth Largest Integer in the Array (compared as strings)Medium