Python · medium · Asked at Google
Asked in ML Engineer interviews, in the Python round.
Use a dict for counts + heapq.nlargest, or Counter.most_common(k).
Count, then select the top k. The interesting part is doing the selection in better than a full sort.
def top_k(items, k):
counts = {}
for x in items:
counts[x] = counts.get(x, 0) + 1
# bucket by count: bucket[c] holds the values 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 outCounting is O(n). Sorting the distinct values by count is O(m log m) for m distinct values; the bucket approach above is O(n) because a count can never exceed n. A heap of size k (heapq.nlargest(k, counts, key=counts.get)) is O(m log k), which is the usual "good" answer and the one to give if buckets feel like overkill.
Say how ties are handled: the bucket version returns whichever tied values were inserted first, which is arbitrary. If the interviewer wants determinism, sort within a bucket.
sorted(counts.items(), key=...)[:k]. Correct, O(m log m), and it reads as not knowing the heap.