Problem 3 · Arrays & Hashing

Top K Frequent Elements

MediumClassic DSA & systems

01Problem

Given an integer array nums and integer k, return the k most frequent elements in any order.

02Key insight

Count with a hash map, then bucket by frequency (frequency is bounded by n) for O(n), avoiding the O(n log k) heap. Mention the heap as the streaming-friendly alternative.

03Solution

solution.pypython14 lines

from collections import Counter

def top_k_frequent(nums: list[int], k: int) -> list[int]:
    counts = Counter(nums)
    buckets: list[list[int]] = [[] for _ in range(len(nums) + 1)]
    for value, freq in counts.items():
        buckets[freq].append(value)
    result: list[int] = []
    for freq in range(len(buckets) - 1, 0, -1):
        for value in buckets[freq]:
            result.append(value)
            if len(result) == k:
                return result
    return result

Try it first.

TimeO(n)

SpaceO(n)

tests.py · 3 passing

tests.pypython3 lines

assert sorted(top_k_frequent([1,1,1,2,2,3], 2)) == [1, 2]
assert top_k_frequent([1], 1) == [1]
assert sorted(top_k_frequent([4,4,5,5,6], 2)) == [4, 5]

04Staff-level follow-up

Top-k heavy hitters on an unbounded stream? Count-Min Sketch plus a size-k min-heap (or Space-Saving / Misra-Gries) gives bounded memory with tunable error.