Heap & Top-K: Beginner

Python's heapq module stores a heap inside an ordinary list. Every function takes the list as its first argument, and the order of the list's elements is the structure. The rule those functions maintain is the heap invariant: for every index i, heap[i] <= heap[2*i + 1] holds, and so does heap[i] <= heap[2*i + 2], wherever those children exist. That rule guarantees the smallest value sits at heap[0] and promises nothing about how the remaining values are ordered.

Read the list as a binary tree stored level by level. Index 0 is the root, the children of index 1 sit at 3 and 4, and the children of index 2 sit at 5 and 6. A value is only ever compared with its own parent and children, so a valid heap can print in an order that looks unsorted. Run heapify on [5, 3, 8, 1, 9, 2], take the minimum off, then push a new value:

heapify rearranges the list in place to [1, 3, 2, 5, 9, 8] and returns None. The 2 at index 2 comes before the 5 at index 3, which is legal because index 3 is a child of index 1. heappop returns the 1. To fill the gap it puts the last element, the 8, at the root, and a sift-down swaps it past its smaller child, leaving [2, 3, 8, 5, 9]. heappush appends the 4 and sifts it up past its parent 8, giving [2, 3, 4, 5, 9, 8]. Each of those 2 calls walks a single root-to-leaf path, which is why each costs O(log n).

Building with heapify costs O(n), less than the O(n log n) of n separate pushes, because it sifts down from the last parent backward. Half the positions are leaves and never move, a quarter can drop at most 1 level, an eighth at most 2, and so on. That series adds up to at most about n level moves for the whole build.

The 2 combined calls have similar names and different results. heappushpop pushes first, so a new item smaller than the root comes straight back and the heap is untouched. heapreplace pops first, so the old root always leaves, even when the new item is smaller.

Either call can be the step of a top-K loop once the heap holds k items: heapreplace behind an if x > h[0] guard, or heappushpop alone, which runs that comparison internally. The bug to avoid is heapreplace with no guard, which evicts the root even when the new value is worse.

heapq only builds min-heaps. To pop the largest value first, store every number negated, so the largest original becomes the smallest stored value, then negate again whenever a value comes out.

Popping an empty heap raises IndexError: index out of range, the message from the C implementation of heappop. Reading h[0] on an empty list raises the list's own wording, IndexError: list index out of range. The second failure raises nothing at all: calling heappush or heappop on a list that was never heapified quietly returns the wrong value.

Every top-K question has the same shape: n values go in, the k best come out, and k is far smaller than n. Sorting answers it in O(n log n) time while holding all n values. The heap answer runs in O(n log k) time and holds k values. The idea behind it reads backwards at first: to keep the k largest, use a min-heap.

Watch the root. After 4 and 1 arrive, then 7, the root is 1, and the value 3 evicts it. Next, 8 evicts the 3, the 2 loses to a root of 4, and 9 evicts the 4. The surviving heap [7, 8, 9] happens to be ascending, but heap order only promises the root, so the function sorts the survivors before returning [9, 8, 7].

The heap's worst case arrives when values come in ascending order, so every one evicts the root. On shuffled input, most values lose their first comparison and cost a single step each.

heapq.nlargest(k, iterable, key=None) and heapq.nsmallest(k, iterable, key=None) implement this pattern and return their results already sorted. The Python documentation adds 2 rules of thumb: max() and min() are more efficient when k = 1, and sorted() is more efficient for larger k. In an interview, write the loop to show the mechanism, then name nlargest as the call you would use in production.

Edge cases are where hand-written versions fail. With k = 0, the loop in the animation skips the push branch, compares the first value against heap[0] on an empty list and raises IndexError: list index out of range. With k larger than the input, the heap never fills and every value comes back. A guard brings the hand-written version in line with the library:

When the ranking value is computed, pass it as key instead of building tuples by hand. For the k points closest to the origin, the key is the squared distance, which ranks points in the same order as distance and skips the square root.

The Kth largest question asks for a single value: the element at position k if the list were sorted in descending order. Duplicates occupy separate positions, so in [3, 2, 3, 1, 2, 4, 5, 5, 6] the 4th largest is 4, reached through 6, 5, 5, 4. Ask about duplicates in the first minute, because the 4th largest distinct value is a different question, and its answer is 3.

The top-K heap from the previous section already holds this answer. Once every value has passed through a min-heap capped at k, the root is the smallest of the k largest values, which is the Kth largest by definition. Seeding the heap with heapify on the first k values replaces k pushes with a single O(k) build.

Quickselect borrows the partition step from quicksort. Pick a pivot, move every smaller value to its left, and drop the pivot into the slot it would occupy in ascending order. If that slot is index n - k, the pivot is the answer. Otherwise the answer lies on a single side, so the loop discards the other side and partitions again, touching a shrinking range each pass.

The target index is 6 - 2 = 4. The first pass uses pivot 4, which lands at index 3, so the answer lies to the right and lo moves to 4. The second pass partitions only indexes 4 and 5: pivot 5 lands at index 4, which is the target, and the function returns 5. The list data came back reordered as [3, 2, 1, 4, 5, 6], because quickselect mutates its input; pass nums[:] when the caller still needs the original order.

Choosing the last element as pivot keeps the animation repeatable and creates a weakness. On input already sorted ascending, the last element is always the largest in range, so each pass removes a single element whenever the target lies to its left, and the work grows to O(n^2). Swapping a random element into position hi before each pass, chosen with random.randint(lo, hi), makes that pattern vanishingly unlikely and bounds the expected cost at about 3.4n comparisons for any k.

The heap's O(k) is all the memory it ever needs. Quickselect's O(1) assumes all n values are already loaded, which a stream never provides.

Put numbers on it. With n = 1,000,000 and k = 100, the heap costs about n log2 k, roughly 6.6 million comparisons in the worst case, and holds 100 values. Quickselect expects at most about 3.4 million comparisons but needs all 1,000,000 values in memory and permission to reorder them. Values already sitting in a Python list on one machine favor quickselect. Values arriving from a stream, or a k tiny next to n, favor the heap. For a single query on a short list, sorted(nums, reverse=True)[k - 1] is the readable choice and a fine answer.

Real heap entries carry data: a priority plus the job or row it belongs to. The usual encoding is a tuple, (priority, payload), which works because tuples compare position by position. Python compares the first elements, and only when they are equal does it move on to the second. That second comparison is where tuple heaps break.

Push 2 jobs that share a priority and carry dict payloads. The first push compares nothing. The second push sifts the new entry up toward its parent, the priorities tie at 2, and Python moves on to compare the 2 dicts with <.

The crash depends on the data. A test suite with distinct priorities never reaches the payload comparison, and the job fails in production on the first day 2 entries share a priority. Equal payloads hide it further: if both dicts held the same keys and values, the tuples would compare equal at that position and no error would appear. Dicts fail this way, and so do objects without __lt__ and a None sitting next to a number.

Put a unique, comparable value between the priority and the payload. A counter from itertools.count() is the standard choice. Priorities decide first, the counter decides every tie, and since no 2 entries share a counter value, the comparison never reaches the payload.

When entries grow past 3 fields, the documentation's other suggestion is a dataclass. With order=True, instances compare as tuples of their fields, and field(compare=False) drops the payload from that tuple.

Sometimes the tie-break is part of the spec: return the most frequent event types with the highest count first, breaking ties alphabetically. Negate the count and keep the name as the second element. In a min-heap, (-5, 'view') comes out before (-3, 'click'), and (-3, 'click') before (-3, 'share'), because 'click' sorts before 'share'.

The pushes produce [(-5, 'view'), (-3, 'click'), (-3, 'share'), (-1, 'like')], and the 3 pops return view, then click, then share. Negation flips the count order while the string keeps its natural order, which is the exact combination the spec asks for.

The heap above held every key. To keep only k, the root must be the weakest kept entry, and under this spec the weakest is the lowest count with the alphabetically last name. Plain (count, word) tuples in a min-heap order counts correctly and names backwards. With k = 2:

When share arrives, (3, 'share') > (3, 'click') is true, so click is evicted even though it wins the tie. A string has no negative, so there is no one-character fix. Hand the full ordering to nsmallest as a key instead; the documentation defines it as equivalent to sorted(iterable, key=key)[:n], ties included:

In a pipeline, top-K rarely arrives as a tidy list of numbers. It arrives as events that must be summed by key or ranked within groups, and often the ranking runs on many partitions whose results have to be merged. Each shape reuses the capped heap and adds a correctness rule of its own.

Customers c3 and c4 tie at 300.0 for second place. The key, negated total then customer ID, hands the place to c3, so a rerun that reads the orders in a different sequence prints the same answer. With heapq.nlargest(2, revenue.items(), key=lambda kv: kv[1]), the tie between c3 and c4 falls back to the dict's insertion order. That order follows the sequence the orders were read in, so 2 runs over the same orders can publish different rankings.

The capped heap assumes each key arrives once with its final value. In a running aggregation, c2 would enter at 45.5 and climb to 255.5 when its second order lands, while the heap keeps the stale 45.5 entry with no way to find and update it. For a batch, rank once after the sums finish. For an hourly refresh, keep running totals in the dict and rebuild the top-K from it with nsmallest each hour, at O(m log k) per refresh.

Finding the 2 slowest queries in each region takes 1 heap per group. A dict maps each region to its own list, setdefault creates that list the first time a region appears, and every event runs the push-or-replace step against its region's heap. Memory is k entries per group, so 50 regions at k = 2 hold 100 entries however long the log runs.

The eu heap rejects q6 at 60 ms because its root is already 120. The us heap fills at q3. After that, q4 evicts the root at the time, which is 90, and q7 evicts the next root, 340. Iterating over sorted(slowest) keeps the region order fixed from run to run.