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

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

Key points

FAQ

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

Practice this for real