← Heaps

Drill #59 — Top K frequent elements

Medium Heap + Map Heaps

In plain English: Find the k numbers that appear most often in an array.

Count frequencies with a hashmap, then use a min-heap of size k to track the top-k — pushing and popping keeps it efficient.

Prompt

Given an array of integers and k, return the k most frequent elements.

Try to write it from scratch before scrolling down.

Solution

import heapq
from collections import Counter

def top_k_frequent(nums, k):
    freq = Counter(nums)
    # nlargest uses a min-heap of size k internally
    return heapq.nlargest(k, freq.keys(), key=freq.get)

# Manual heap approach:
def top_k_frequent_manual(nums, k):
    freq = Counter(nums)
    heap = []
    for num, cnt in freq.items():
        heapq.heappush(heap, (cnt, num))
        if len(heap) > k:
            heapq.heappop(heap)  # remove least frequent
    return [num for cnt, num in heap]

# Test: top_k_frequent([1,1,1,2,2,3], 2) == [1, 2]
O(n log k) time · O(n) space

Related Micro Drills

← Drill #58 Drill #60 →