Top K Frequent Elements

Problem

Given an integer array nums and an integer k, return the k most frequent elements within the array. The test cases are generated such that the answer is always unique. You may return the output in any order.

Example 1

Input:  nums = [1, 2, 2, 3, 3, 3], k = 2
Output: [2, 3]

Example 2

Input:  nums = [7, 7], k = 1
Output: [7]

Constraints

1 <= len(nums) <= 10^4
-1000 <= nums[i] <= 1000
1 <= k <= number of distinct elements in nums

Thinking

Let $n$ = len(nums), $u$ = number of distinct values ( $u \le n$ , and here $u \le 2001$ ).

Brute force idea: count every value, then sort the $u$ distinct values by count and take the first $k$ : $O(n + u \log u)$ . That’s already fine; the follow-up asks for better than $O(n \log n)$ .
Key insight: step 1 is always count with a hash map ( $O(n)$ $O (n)$ ). Step 2 is “pick the top $k$ $k$ by count”, and there are three ways to do it:
- heap of size $k$ : $O(u \log k)$ ;
- bucket sort: a count can only be between 1 and $n$ , so put each value in buckets[count] and walk from the highest count down: $O(n)$ ;
- full sort: $O(u \log u)$ .
Unique answer guarantee: the bucket that takes you past $k$ can never have a tie, so you can add whole buckets with extend and stop when len(result) >= k.
Edge cases to remember: one element ([7]), every element the same (count $= n$ , so the bucket list needs index $n$ ), $k$ = number of distinct values, negative numbers.

Solution

Approach A — my version (bucket sort, after fixing the two bugs)

Complexity: Time $O(n)$ · Space $O(n)$ . The Big-O is already optimal. It’s slow in practice because of constant factors:

n + 1 buckets, sized by len(nums). With $n = 10^4$ that is 10,001 Python list objects (~700 KiB), built in a Python loop, even though the highest count is often around 5. This is the main memory cost and about a quarter of the runtime.
The bucket walk never stops. After result has $k$ elements the if keeps failing, but the loop still visits all 10,001 buckets. temp[::-1] also copies the whole list first.
Counting in a Python loop. frequency[num] += 1 runs $n$ times in bytecode; Counter(nums) does the same counting in C, about 2× faster.
int(x) is left over from the old str keys and does nothing now.

The code is unchanged except for the signature and comments.

from collections import Counter, defaultdict
import heapq

def top_k_frequent_mine(nums: list[int], k: int) -> list[int]:
    frequency = defaultdict(int)
    temp = [[] for _ in range(len(nums) + 1)]  # SLOW: n+1 lists, most of them stay empty
    for num in nums:
        frequency[num] += 1                      # SLOW: counting loop in Python, Counter is C
    for x, y in frequency.items():
        temp[y].append(int(x))                   # int() does nothing now
    result = []
    for z in temp[::-1]:                         # copies the list, then visits every bucket
        if len(result) < k:
            result.extend(z)                     # no break: keeps looping after it has k
    return result

Approach B — bucket sort, done efficiently (the $O(n)$ interview answer)

Complexity: Time $O(n)$ · Space $O(n)$ worst case, but only $O(u + \text{max count})$ in practice. Same idea as mine with three fixes: count with Counter, make only max count + 1 buckets, and break once there are $k$ results. This alone makes it ~3× faster and cuts peak memory ~7×.

def top_k_frequent_bucket(nums: list[int], k: int) -> list[int]:
    count = Counter(nums)
    buckets = [[] for _ in range(max(count.values()) + 1)]  # index = count
    for num, c in count.items():
        buckets[c].append(num)
    result = []
    for bucket in reversed(buckets):  # highest count first, no copy
        result.extend(bucket)
        if len(result) >= k:
            break
    return result

Approach C — heap of size $k$ (`heapq.nlargest`)

Complexity: Time $O(n + u \log k)$ · Space $O(u + k)$ . Count, then keep only the $k$ largest counts in a min-heap. nlargest does this in C. It’s the answer to give when the data is a stream or $k$ is much smaller than $u$ .

def top_k_frequent_heap(nums: list[int], k: int) -> list[int]:
    count = Counter(nums)
    return heapq.nlargest(k, count, key=count.get)

Approach D — `Counter.most_common` ✅ best practice (Pythonic)

Complexity: Time $O(n + u \log k)$ · Space $O(u + k)$ . most_common(k) calls heapq.nlargest internally, so this is Approach C in one line. In everyday Python, write this. In an interview, write it, then explain B (the $O(n)$ bucket sort) when asked for better than $O(n \log n)$ .

def top_k_frequent(nums: list[int], k: int) -> list[int]:
    return [num for num, _ in Counter(nums).most_common(k)]

Scratch / Experiments

Worst-ish case at the max size: $10^4$ random values from the full range $[-1000, 1000]$ , so $u \approx 2001$ and the highest count is only ~15.

All four are $O(n)$ or close to it, so the gaps are constant factors:

A vs B: same algorithm. B is ~3× faster only because it counts in C, builds ~15 buckets instead of 10,001, and stops early.
B vs C/D: bucket sort is $O(n)$ and the heap is $O(n + u \log k)$ , but with $k$ small, $\log k$ is tiny and nlargest runs in C, so they finish about even. Big-O tells you how it grows, not which is faster at $n = 10^4$ .
The second block measures peak memory: A’s 10,001 empty lists are almost all of it.

LeetCode’s runtime/memory percentiles are noisy (±20% between runs of the same code), so treat them as a hint, not a verdict.

import random, tracemalloc

random.seed(0)
nums = [random.randint(-1000, 1000) for _ in range(10_000)]
k = 10

for fn in APPROACHES:
    print(fn.__name__)
    %timeit fn(nums, k)

for fn in APPROACHES:
    tracemalloc.start()
    fn(nums, k)
    peak = tracemalloc.get_traced_memory()[1]
    tracemalloc.stop()
    print(f"{fn.__name__:24s} peak {peak / 1024:6.0f} KiB")

Takeaways

Pattern: “top $k$ / $k$ most frequent / $k$ largest” → count with a hash map, then select: Counter.most_common(k) / heapq.nlargest ( $O(u \log k)$ ), or bucket sort by count when the counts are bounded by $n$ ( $O(n)$ ).
Mistakes I made:
- Sized the bucket list as len(nums): a value can appear len(nums) times, so I needed len(nums) + 1 (crashed on [7, 7]).
- Used <= k as “still need more”: that’s < k. Off-by-one in a stop condition: trace the example by hand and check the last iteration.
- Used str(num) as the dict key: ints are hashable, so convert nothing.
- Optimal Big-O can still be slow. Mine was $O(n)$ but allocated 10,001 lists and never stopped looping. Size buffers by what’s needed (max(count.values())), break when done, and prefer C-implemented builtins (Counter, heapq) over hand-written loops.
Related problems: Kth Largest Element in an Array (#215, heap / quickselect for top $k$ ), Top K Frequent Words (#692, same idea plus a tie-break by word), Sort Characters By Frequency (#451, count + bucket sort), K Closest Points to Origin (#973, heapq.nsmallest with a key), Group Anagrams (#49, the counting-with-a-dict step).