Top K Frequent Elements (LeetCode 347)

The problem

You get an integer array nums and an integer k. Return the k values that appear most often. The answer is unique, and it may be returned in any order.

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

1 appears three times, 2 twice and 3 once. The two most frequent are 1 and 2.

How to think about it

Start with brute force. For each distinct value, count it by scanning the array, then sort the values by count. Counting that way is O(n²). The first fix is obvious once you’ve done this lesson: count everything in one pass with a hash map.

Counting is half the job. After one pass you have a map like {1: 3, 2: 2, 3: 1}. Now the problem is “pick the k keys with the largest values”. Sorting all u distinct keys by count is O(u log u), which is fine but does more work than needed when k is small.

The key insight: you only need the top k, so keep only k. A heap of size k gives the k largest counts in O(u log k). Python’s Counter.most_common(k) and heapq.nlargest(k, …) both do exactly this under the hood. For a guaranteed O(n), use bucket sort: make a list of buckets indexed by count (0 to n), drop each value into its count’s bucket, and read buckets from the top until you have k values.

Which pattern is this? Hash map frequency counting, followed by top-k selection. The tell is “the k most / least frequent”.

Approach

1. Count every value with Counter(nums).
2. Ask for the k most common entries. most_common(k) uses a heap internally when k is given.
3. Return just the values from those (value, count) pairs.

Solution

from collections import Counter

def top_k_frequent(nums, k):
    cnt = Counter(nums)                          # value -> frequency
    return [x for x, _ in cnt.most_common(k)]    # k highest counts
# top_k_frequent([1,1,1,2,2,3], 2) -> [1,2]

Alternative: explicit heap

import heapq
from collections import Counter
def top_k_frequent(nums, k):
    cnt = Counter(nums)
    return sorted(heapq.nlargest(k, cnt.keys(), key=cnt.get))   # k keys with largest counts
# top_k_frequent([1,1,1,2,2,3], 2) -> [1,2]

This is the same algorithm with the heap made visible, which is how you’d write it in a language without Counter (and it’s how Lesson 14 frames the problem). The final sorted orders the answer by value, not frequency, which is harmless because any order is accepted. Choose this version if the interviewer wants to see the heap.

Dry run

On nums = [1, 1, 1, 2, 2, 3], k = 2:

With nums = [1] and k = 1, the counter is {1: 1} and the answer is [1].

Complexity

• Time: O(n log k). Counting is O(n). Selecting the top k from u ≤ n distinct values with a size-k heap is O(u log k).
• Space: O(n). The counter holds up to n distinct values; the heap holds k.

Edge cases and pitfalls

• k equals the number of distinct values: every value is returned.
• Ties: the problem guarantees a unique answer, so you don’t need a tie-breaking rule. If asked, say how you’d break ties (for example, by smaller value).
• Sorting the whole array by frequency with a custom key works but is O(n log n) and builds a much bigger list than needed.
• Min-heap vs max-heap: to keep the k largest in a size-k heap by hand, use a min-heap and pop whenever it grows past k.
• Return values, not counts: most_common returns pairs; strip the counts.

Related problems

• Kth Largest Element in an Array (LeetCode 215): the same size-k heap, on values instead of counts.
• Sort Characters By Frequency (LeetCode 451): count, then order everything by frequency.
• Group Anagrams (LeetCode 49): another hash-map bucketing problem.