Top K Frequent Elements: Bucket Sort in O(n)

"Return the k most frequent elements" has a standard answer, a heap, and a follow-up that catches people out: "your solution is O(n log k); can you do better?" You can. The counts in this problem are not arbitrary numbers. They are small integers that can never exceed the length of the input, and a small bounded integer is an array index. Once you see that, the ranking step turns into plain indexing and the whole solution is linear.

The problem it solves

Given nums and k, return the k values that appear most often. For [1, 1, 1, 2, 2, 3] and k = 2, the answer is [1, 2]. The order of the answer usually does not matter, and the problem normally promises that the answer is unique, so there is no tie at the k-th place to break.

Every solution starts the same way: count occurrences with a hash map in one pass. The question is how to pick the top k from those counts.

• Sort the counts: O(u log u) for u distinct values, which is O(n log n) in the worst case.
• Heap of size k: push each value with its count and pop whenever the heap grows past k, for O(u log k). This is the approach from top k elements with a heap.
• Buckets: O(n), with no comparisons between values at all.

The intuition

Sorting by count feels necessary because we want the values in count order. But sorting is expensive only when the keys can be anything. Here every count lies between 1 and n, because a value cannot appear more often than the array is long.

So make an array of n + 1 buckets, where buckets[c] is the list of values that appear exactly c times. Every value goes into exactly one bucket, found by direct indexing. The buckets are then already in count order by position. Walk them from index n down to 1, collecting values, and stop as soon as you have k.

This is bucket sort with a guarantee that makes it safe: the key range is bounded by the input size, so the bucket array never costs more than the input. The same idea is behind counting sort, which indexes by the value instead of by the count.

Top-k is one of the patterns collected on the patterns cheat sheet, with the heap and bucket versions side by side.

Watch it run

The animation uses the classic input with k = 2. It opens with the key fact: a heap would give the top 2 in O(n log k), but no count can be larger than the array itself. First the tally, one pass with a map: 1 three times, 2 twice, 3 once. Then each value drops into the bucket named by its count. Counts are small integers, so they are addresses. The walk starts from the busiest bucket: bucket 3 holds 1. Bucket 2 hands over 2, which makes two values, so the walk stops there and bucket 1 is never read. The last frame names the three phases: tally, scatter, walk, three linear passes, and not one value was compared with another.

The code

The lesson's function with the three passes labelled, on inputs that include strings and a single element:

def top_k(nums, k):
    counts = {}
    for x in nums:                                 # pass 1: tally
        counts[x] = counts.get(x, 0) + 1
    buckets = [[] for _ in range(len(nums) + 1)]   # index = how often
    for value, c in counts.items():                # pass 2: scatter
        buckets[c].append(value)
    out = []
    for c in range(len(nums), 0, -1):              # pass 3: walk down
        out += buckets[c]
        if len(out) >= k:
            return out[:k]
    return out

print(top_k([1, 1, 1, 2, 2, 3], 2))              # [1, 2]
print(top_k(["b", "a", "b", "c", "a", "b"], 1))  # ['b']
print(top_k([4, 4, 5, 5, 6], 2))                 # [4, 5]
print(top_k([7], 1))                             # [7]

For comparison, the heap version in two lines. heapq.nlargest keeps a heap of size k internally:

from collections import Counter
import heapq

def top_k_heap(nums, k):
    counts = Counter(nums)
    return heapq.nlargest(k, counts, key=counts.get)   # O(n log k)

print(top_k_heap([1, 1, 1, 2, 2, 3], 2))         # [1, 2]

Both are checked against a brute force that sorts every distinct value by count, on 3,000 random inputs. Ties at the cut-off can legitimately be broken differently, so the check compares the multiset of counts returned, not the values:

import random

def brute(nums, k):
    counts = Counter(nums)
    ranked = sorted(counts, key=lambda v: -counts[v])
    return counts, ranked[:k]

random.seed(22)
ok = True
for _ in range(3000):
    nums = [random.randint(0, 9) for _ in range(random.randint(1, 30))]
    k = random.randint(1, len(set(nums)))
    counts, want = brute(nums, k)
    got = top_k(nums, k)
    # ties may be broken differently, so compare the counts, not the values
    ok &= len(got) == k == len(set(got))
    ok &= sorted(counts[v] for v in got) == sorted(counts[v] for v in want)
    ok &= sorted(counts[v] for v in top_k_heap(nums, k)) == sorted(counts[v] for v in want)
print(ok)                                        # True

The complexity

• Time: O(n). The tally is one pass over nums, the scatter is one pass over at most n distinct values, and the walk visits at most n buckets.
• Space: O(n) for the map and the n + 1 buckets. The heap version needs only O(u + k), which matters when there are few distinct values.
• In practice: the heap is often just as fast. When k is small, log k is a tiny constant, and the bucket array allocates n + 1 lists even if only a handful are used.

Where it goes wrong

• Allocating n buckets instead of n + 1. When every element is the same value, its count is n, and index n must exist.
• Sizing the buckets by the number of distinct values. Counts are bounded by len(nums), not by how many different values there are.
• Walking up instead of down. Starting at bucket 1 finds the least frequent values first.
• Forgetting to slice. A bucket can hold several values; if taking the whole bucket overshoots k, trim to out[:k].
• Using this for a stream. Buckets need the final counts. For a stream that never ends, a heap or a sketch is the tool; the heap version is covered in top k frequent in a stream.

When it shows up in interviews

"Top k frequent elements" is one of the most common hash map questions, and the bucket answer is what the "better than O(n log k)" follow-up is looking for. Close relatives are "top k frequent words", where ties must be broken alphabetically and a heap with a compound key is simpler, and "sort characters by frequency", where the bucket walk produces the whole ordering. Counting with a map is the same first step as in valid anagram and group anagrams.

How to say it in an interview

"First I count occurrences with a hash map. A heap of size k would give O(n log k), but the counts are integers between 1 and n, so I can use them as indexes. I make n + 1 buckets, put each value into the bucket for its count, then walk from bucket n downwards collecting values until I have k. Every step is a linear pass with no comparisons, so it is O(n) time and O(n) space. If k is tiny or the input is a stream, I would use the heap instead."