~/problems / Heaps / Heaps and priority queues

Top K Frequent Elements

easy ~20 min

Write top_k_frequent(nums, k) that returns the k distinct values that occur most often in nums, in any order. The inputs are chosen so the answer is unique (there's never a tie at the cut-off).

top_k_frequent([4, 4, 4, 9, 9, 2], 2)   # [4, 9]  (or [9, 4])
top_k_frequent([-1, -1, 3], 1)          # [-1]
top_k_frequent([8], 1)                  # [8]

Constraints: 1 <= k <= number of distinct values; len(nums) up to 200,000 with up to 100,000 distinct values.

Aim for better than O(n log n), e.g. O(n log k) or O(n). Calling nums.count(x) for each distinct x is O(n²) and too slow.

Show hint

Count each value once. To pick the k biggest counts without sorting all of them, keep only the best k seen so far in a structure that makes the weakest of them cheap to find, or group values by their count.

Topic: Heaps and priority queues. heapq: top-k, k-way merge, two heaps for a running median.

Read the visual guide
0:00
Ctrl ' run · Ctrl ↵ submit
esc