Interview challenge · medium · array, hash-map, heap, bucket-sort
Top K Frequent Elements
The challenge
Given an integer array nums and an integer k, return the k most frequent elements, most frequent first.
Example
[1, 1, 1, 2, 2, 3], k = 2 -> [1, 2]
[4, 4, 5, 5, 5, 6], k = 1 -> [5]Idea: count how many times each number appears. A number can appear at most n times, so put each number in bucket number *count* and walk the buckets from n down to 1 until you have k numbers: O(n) with no sorting. The C version keeps the counts in a small table and picks the largest k times.
The solution
# Top K Frequent Elements - ayotl.dev
def top_k(nums, k):
count = {}
for x in nums:
count[x] = count.get(x, 0) + 1
n = len(nums)
# bucket c: numbers seen c times
buckets = [[] for _ in range(n + 1)]
for x, c in count.items():
buckets[c].append(x)
out = []
for c in range(n, 0, -1):
for x in buckets[c]:
out.append(x)
if len(out) == k:
return out
return out
print(top_k([1, 1, 1, 2, 2, 3], 2))
print(top_k([4, 4, 5, 5, 5, 6], 1))
$ python3 top_k_frequent_elements.py
[1, 2]
[5]