Top-K problems

Top-K problems ask for the k largest, smallest, most frequent or closest items, or for just the k-th one. Sorting everything works in O(n log n), but it does more than needed and needs all the data in memory.

The heap trick is counter-intuitive: for the k largest you use a min-heap, and for the k smallest a max-heap. The heap holds the current best k, and its top is the worst of them, the one that would be kicked out next. Each new item is compared with that top in O(1), and only winners pay O(log k) to enter. The heap never grows beyond k, so it also works on a stream of unknown length.

1. 1Decide the direction. k largest: min-heap (the top is the smallest of the winners). k smallest or k closest: max-heap (in Python, push negated keys).
2. 2For each item: if the heap has fewer than k items, push it.
3. 3Otherwise compare the item with the top. If it is better than the top, replace the top with it (heapreplace pops and pushes in one O(log k) step). If not, skip it.
4. 4At the end the heap holds exactly the top k items in no particular order; sort them if order is needed (O(k log k)). The k-th largest is simply the top of the min-heap.
5. 5For 'most frequent', first count with a hash map, then run the same size-k heap over (count, value) pairs.

import heapq
from collections import Counter

def k_largest(nums, k):
    heap = []                              # min-heap holding the k biggest so far
    for x in nums:
        if len(heap) < k:
            heapq.heappush(heap, x)
        elif x > heap[0]:                  # beats the weakest of the current top k
            heapq.heapreplace(heap, x)     # pop the smallest, push x
    return sorted(heap, reverse=True)      # heap[0] alone is the k-th largest

def top_k_frequent(nums, k):
    heap = []
    for val, count in Counter(nums).items():
        heapq.heappush(heap, (count, val))
        if len(heap) > k:
            heapq.heappop(heap)            # drop the least frequent
    return [val for count, val in sorted(heap, reverse=True)]

k_largest keeps a min-heap of the best k and replaces the top only when a new item beats it. top_k_frequent counts first, then keeps the k most frequent (count, value) pairs by popping the smallest whenever the heap exceeds k. Both are O(n log k) time and O(k) heap memory (plus O(u) for the counter over u distinct values).

Cost: Size-k heap: O(n log k) time, O(k) space. Sorting: O(n log n). Quickselect: O(n) average, O(n^2) worst. heapq.nlargest(k, data) uses the same size-k heap idea.