Top K, Merging Sorted Lists and Max Heaps
28 min · Intermediate
Need help? Ask a tutor on WhatsApp (opens in a new tab)Heaps solve three common jobs: the best k of many values, merging sorted lists, and a max heap in a language that only has a min-heap. Each uses the heapq basics you already know, with one small trick.
A max heap with negated values
A max heap by negation stores -x for every value x in a min-heap, so the top is the largest x. Negate on the way in and on the way out. It works for numbers only.
heapq only has a min-heap. To get the largest value first, store each value with its sign flipped. The smallest negative number is the largest original number. Flip the sign again when you read it.
import heapq
scores = [40, 90, 70]heap = [-s for s in scores]heapq.heapify(heap)print(-heap[0])print(-heapq.heappop(heap))Negation costs nothing and keeps all heap operations O(log n). It works for numbers only. For strings or objects, use a tuple key instead.
Top k: a small heap of the best so far
Top-k with a heap keeps a heap of at most k items and drops the worst whenever it is full. It takes O(n log k) time and O(k) space and works on a stream.
To keep the k largest values, hold a min-heap of size k. Push each new value. If the heap grows past k, pop the smallest. What stays are the k largest, and the top is the k-th largest.
import heapq
def top_k(nums, k): heap = [] for n in nums: heapq.heappush(heap, n) if len(heap) > k: heapq.heappop(heap) return sorted(heap, reverse=True)Example 1 · Top 3 scores: sort or heap
import heapq
scores = [55, 91, 62, 78, 91, 40, 85]print(sorted(scores, reverse=True)[:3])print(heapq.nlargest(3, scores))heap = []for s in scores: heapq.heappush(heap, s) if len(heap) > 3: heapq.heappop(heap)print(sorted(heap, reverse=True))[91, 91, 85] [91, 91, 85] [91, 91, 85]
What this code does
- All three give the same answer. Sorting costs O(n log n). The heap versions cost O(n log k) with k = 3.
- The hand-made loop is what an interviewer expects.
nlargestis the short form for real code.
Example 2 · K-th largest in a stream
import heapq
k = 3heap = []for value in [4, 5, 8, 2, 9, 1]: heapq.heappush(heap, value) if len(heap) > k: heapq.heappop(heap) if len(heap) == k: print(heap[0])4 4 5 5
What this code does
- The heap holds the 3 largest values so far. Its top is the smallest of them, which is the third largest.
- Each new value costs O(log k) and the heap never grows past 3.
Brute force sorts everything in O(n log n). The heap does O(n log k) and holds only k values, so it suits a long stream. When the list is small, heapq.nlargest(k, nums) does the same.
Merge k sorted lists
A k-way merge keeps the front item of each of k sorted lists in a heap. Pop the smallest, then push its successor. It takes O(n log k) time for n items in total.
Each list is sorted, so the next smallest overall is the front of one list. Keep one front value per list in a heap. Pop the smallest, then push the next value from the same list.
import heapq
def merge_k(lists): heap = [(lst[0], i, 0) for i, lst in enumerate(lists) if lst] heapq.heapify(heap) result = [] while heap: value, i, j = heapq.heappop(heap) result.append(value) if j + 1 < len(lists[i]): heapq.heappush(heap, (lists[i][j + 1], i, j + 1)) return resultExample 3 · Merge three sorted lists
import heapq
def merge_k(lists): heap = [(lst[0], i, 0) for i, lst in enumerate(lists) if lst] heapq.heapify(heap) result = [] while heap: value, i, j = heapq.heappop(heap) result.append(value) if j + 1 < len(lists[i]): heapq.heappush(heap, (lists[i][j + 1], i, j + 1)) return result
lists = [[1, 4, 9], [2, 3, 10], [5]]print(merge_k(lists))print(list(heapq.merge(*lists)))[1, 2, 3, 4, 5, 9, 10] [1, 2, 3, 4, 5, 9, 10]
What this code does
- The list index
iin the tuple tells us where to take the next value from. heapq.mergeis the library version. It gives the same result and works lazily.
The heap holds at most k entries, so each step costs O(log k). For n values in total the time is O(n log k) and the extra space is O(k). Joining and sorting would cost O(n log n).
Tuples, ties and a counter
A tie-breaking counter is a number placed between the priority and the item in a heap tuple. It makes every tuple unique, so Python never compares the items, and it keeps equal priorities first in, first out.
A heap can hold tuples. Python compares them item by item, so (priority, name) orders by priority and then by name. If two tuples tie and the next item cannot be compared, add a running counter before it.
import heapqfrom itertools import count
tick = count()heap = []for priority, job in [(2, {"id": 1}), (1, {"id": 2}), (2, {"id": 3})]: heapq.heappush(heap, (priority, next(tick), job))print(heapq.heappop(heap)[2])Example 4 · Most frequent words with negated counts
import heapqfrom collections import Counter
words = "a b a c b a d c a b".split()counts = Counter(words)heap = [(-n, word) for word, n in counts.items()]heapq.heapify(heap)for _ in range(2): n, word = heapq.heappop(heap) print(word, -n)a 4 b 3
What this code does
- The tuple (-count, word) puts the biggest count first and breaks ties by the smaller word.
Countercounts in O(n), then each pop costs O(log m) for m different words.
Example 5 · Equal priorities stay in order
import heapqfrom itertools import count
tick = count()heap = []for priority, name in [(2, "print"), (1, "email"), (2, "backup"), (1, "login")]: heapq.heappush(heap, (priority, next(tick), name))while heap: print(heapq.heappop(heap)[2])email login print backup
What this code does
- Priority 1 comes out first. Inside the same priority, the job that arrived first comes first.
- Without the counter, equal priorities would be ordered by name instead.
Dictionaries cannot be compared with <, so a tie on priority would raise a TypeError. The counter breaks the tie first and keeps jobs of equal priority in the order they arrived.
Next step
Check your understanding
10 quick questions