Top-k Problems and Merging

The k biggest, without sorting everything

To find the 3 highest scores out of a million, keep a min-heap of size 3. For each new score, if it beats the smallest of your top 3 (the heap's root), swap it in:

import heapq

def top_k(nums, k):
    heap = []
    for x in nums:
        heapq.heappush(heap, x)
        if len(heap) > k:
            heapq.heappop(heap)       # throw out the smallest
    return sorted(heap, reverse=True)

print(top_k([5, 12, 3, 41, 7, 19, 2, 33], 3))

Output

[41, 33, 19]

Each step works on a heap of only k items: O(n log k), which is much faster than sorting when k is small. Python also has heapq.nlargest(k, nums) that does this for you.

Think of it like this: A talent show judge who only keeps the current top 3 photos on their desk. Every new photo is compared to the worst of the three.

Merging sorted lists

Put the first item of each list in a heap. Repeatedly take the smallest and replace it with the next item from the same list:

import heapq

def merge_lists(lists):
    heap = [(lst[0], i, 0) for i, lst in enumerate(lists) if lst]
    heapq.heapify(heap)
    out = []
    while heap:
        value, which, pos = heapq.heappop(heap)
        out.append(value)
        if pos + 1 < len(lists[which]):
            heapq.heappush(heap, (lists[which][pos + 1], which, pos + 1))
    return out

print(merge_lists([[1, 4, 9], [2, 3, 10], [5, 6]]))

Output

[1, 2, 3, 4, 5, 6, 9, 10]

Remember: Need the smallest/biggest item over and over while things keep changing? Use a heap.