Python · medium · Reported at Google
Find the k most frequent elements in a list (no libraries).
Asked in ML Engineer interviews, in the Python round.
Short answer
Use a dict for counts + heapq.nlargest, or Counter.most_common(k).
How to answer it
Definition
Count how often each value appears, then return the k values with the highest counts. Counting is the easy half. The interesting half is selecting the top k without paying for a full sort.
Approach
- Count occurrences into a dictionary. That is O(n) and unavoidable.
- Bucket the values by their count. A count can never exceed
len(items), so an array indexed by count is bounded and cheap. - Walk the buckets from the highest count downwards, collecting until k values are out.
Example
def top_k(items, k):
counts = {}
for x in items:
counts[x] = counts.get(x, 0) + 1
# bucket[c] holds every value seen exactly c times.
buckets = [[] for _ in range(len(items) + 1)]
for value, c in counts.items():
buckets[c].append(value)
out = []
for c in range(len(buckets) - 1, 0, -1):
for value in buckets[c]:
out.append(value)
if len(out) == k:
return out
return out
Output
>>> top_k(["a", "b", "a", "c", "b", "a"], 2)
['a', 'b']
Complexity
Counting O(n) unavoidable
Bucket selection O(n) the version above
Sorting by count O(m log m) m = number of distinct values
Heap of size k O(m log k) heapq.nlargest, the usual good answer
Bucketing wins because it exploits something specific about the data: counts are bounded by n, so they can be used directly as array indexes. That is the counting-sort trick, and it is worth naming out loud, because it shows you know why it is O(n) rather than having memorised that it is.
The heap alternative
If the bucket idea does not arrive, this is the answer to give, and it still beats sorting everything:
import heapq
def top_k(items, k):
counts = {}
for x in items:
counts[x] = counts.get(x, 0) + 1
return heapq.nlargest(k, counts, key=counts.get)
It is preferable when the number of distinct values is huge and k is small, because it holds only k items at a time.
Edge cases to say out loud
- k is larger than the number of distinct values. Return everything rather than crashing.
- k is zero or negative. Decide the behavior and state it.
- Ties on the boundary. Two values share the k-th count, and which one comes back is arbitrary unless the question defines a tiebreak.
- The list is empty. Return an empty list.
- Values are unhashable, such as lists. A dictionary count will not work at all.
Key points
- Counting is O(n); every improvement lives in the selection step.
- Bucketing is O(n) because a count is bounded by the input length.
heapq.nlargest(k, counts, key=counts.get)is O(m log k) and a perfectly good answer.- Sorting the whole count map is the answer to avoid, and the one most candidates reach for first.
- Mentioning ties and an oversized k before being asked is most of the difference between a pass and a strong pass.
FAQ
- Why not use
collections.Counter? Itsmost_commondoes exactly this. The question forbids libraries because the point is whether you can explain the selection step. - Is bucketing always fastest? No. It allocates an array of size n+1, so with very few distinct values in a very long list the heap is leaner.
- Does the output need to be sorted? Only if the question says so, and sorting k items afterwards is cheap.
The 30-second answer
I would count the values into a dictionary first, which is O(n) and unavoidable. The real question is how to select the top k without sorting everything.
My first answer would be a heap: keep k items and push the counts through it, which is O(m log k) for m distinct values. If they want better than that, I would bucket by count. A count can never exceed the length of the list, so I can index an array by count and walk it downwards, and that makes the whole thing O(n).
Before I finish I would say what happens on ties, and what I return if k is larger than the number of distinct values.
In the interview
- What they are checking: the count-then-select structure and the complexity of each step.
- Common mistake:
sorted(counts.items(), key=...)[:k]. Correct, O(m log m), and it reads as not knowing the heap.