Top K Frequent Elements: Return the k Most Common Values

Problem

Given a list of integers nums and an integer k, return the k values that occur most often, in any order. Assume the answer is unique: no tie straddles the boundary between the k-th and (k+1)-th most frequent value. This is widely known as LeetCode 347, Top K Frequent Elements.

k is between 1 and the number of distinct values. The list can hold up to about 10^5 integers, including negatives.

Examples

nums = [4, 4, 4, 7, 7, 1],  k = 2   ->  [4, 7]
nums = [-2],                k = 1   ->  [-2]
nums = [5, 3, 5, 3, 5, 9],  k = 1   ->  [5]

Approach 1: brute force

Count with a dictionary, sort the distinct values by count descending, take the first k.

from collections import Counter

def top_k_sort(nums, k):
    counts = Counter(nums)
    return sorted(counts, key=counts.get, reverse=True)[:k]

Complexity: O(n + d log d) time for d distinct values (O(n log n) in the worst case), O(n) space. This is a good first answer; interviewers then ask you to beat O(n log n).

Approach 2: optimal (bucket sort by frequency)

Key insight: a frequency is an integer between 1 and n, so you can index values by their count instead of sorting.

Walkthrough on [4, 4, 4, 7, 7, 1], k = 2:

1. Counts: {4: 3, 7: 2, 1: 1}.
2. Buckets (index = count): [[], [1], [7], [4], [], [], []].
3. Read from index 6 down: index 3 gives 4, index 2 gives 7. That is two values, stop.

def top_k_frequent(nums, k):
    counts = {}
    for num in nums:
        counts[num] = counts.get(num, 0) + 1
    buckets = [[] for _ in range(len(nums) + 1)]
    for value, c in counts.items():
        buckets[c].append(value)
    result = []
    for c in range(len(buckets) - 1, 0, -1):
        for value in buckets[c]:
            result.append(value)
            if len(result) == k:
                return result
    return result

Complexity: O(n) time, O(n) extra space.

Approach 3: min-heap of size k

Push (count, value) pairs onto a min-heap and pop whenever it grows past k. The heap always holds the k largest counts seen so far. This is what heapq.nlargest does internally.

import heapq

def top_k_heap(nums, k):
    counts = Counter(nums)
    heap = []
    for value, c in counts.items():
        heapq.heappush(heap, (c, value))
        if len(heap) > k:
            heapq.heappop(heap)
    return [value for c, value in heap]

Complexity: O(n + d log k) time, O(d + k) space. Prefer it when k is much smaller than d, or when counts arrive as a stream and you cannot allocate n buckets.

Tests

import random

for f in (top_k_frequent, top_k_heap, top_k_sort):
    assert sorted(f([4, 4, 4, 7, 7, 1], 2)) == [4, 7]
    assert f([-2], 1) == [-2]                              # single element
    assert f([5, 3, 5, 3, 5, 9], 1) == [5]
    assert sorted(f([1, 2, 3], 3)) == [1, 2, 3]            # k equals distinct count
    assert sorted(f([-1, -1, -3, -3, -3, 0], 2)) == [-3, -1]   # negatives
    assert f([10**9] * 4 + [1], 1) == [10**9]              # large values

big = [i % 1000 for i in range(100_000)] + [7] * 50
assert top_k_frequent(big, 1) == [7]

random.seed(5)
for _ in range(300):
    arr = [random.randint(-5, 5) for _ in range(random.randint(1, 20))]
    counts = Counter(arr)
    freqs = sorted(counts.values(), reverse=True)
    k = random.randint(1, len(freqs))
    if k < len(freqs) and freqs[k - 1] == freqs[k]:
        continue                                           # answer not unique: skip
    expected = sorted(top_k_sort(arr, k))
    assert sorted(top_k_frequent(arr, k)) == expected == sorted(top_k_heap(arr, k))

Edge cases and pitfalls

• The bucket list needs n + 1 slots because a single value can occur n times.
• Python’s heapq is a min-heap. To keep the k largest you pop the smallest, not the largest.
• If ties are possible, ask how to break them (for example by smaller value) and include the tie-breaker in the sort key.
• Counter.most_common(k) is fine to mention, but be ready to explain what it does underneath.

Where this shows up in data engineering

“Top N products by order count” is this problem. In SQL it is GROUP BY plus ORDER BY ... LIMIT or a ROW_NUMBER() filter. On an unbounded stream you cannot keep every count, so approximate structures such as Count-Min Sketch with a heap of heavy hitters take over.