Heap Basics & Top K

Reach for a heap when you need to rank entries on the go and keep fetching the top one (the smallest or the largest). Sorting ranks everything once; a heap keeps the ranking correct while items arrive and leave. The most common use is Top K: keep only the K best items seen so far and throw the rest away.

The idea#

Python's heapq works on a plain list and is a min-heap only: the smallest item always sits at heap[0]. These are the calls you will use most:

Before Python 3.14 there is no public max-heap API, and LeetCode runs older versions. To track the largest values, negate them on the way in and negate them back on the way out.

Top K: keep the weakest link on top#

The heap you need is the opposite of what the question asks for, which feels backward at first:

• K largest → a min-heap of size K. The smallest item in the heap is the "weakest link". When a bigger item arrives, evict the weakest link.
• K smallest → a max-heap (negated values) of size K. Evict the biggest.

Each of the N items costs at most one $O(\log K)$ heap operation, so the total is $O(N \log K)$ time and $O(K)$ space. When K is small this beats sorting ($O(N \log N)$), and it works on a stream you can't hold in memory.

How to spot it#

These cues point to a heap, and to the size-k version in particular.

• "Kth largest / smallest", "top K", "K closest", "K most frequent".
• Repeatedly take the largest (or smallest) item, change it, and put it back (a simulation).
• Data arrives as a stream and you must answer after every insert.
• N is large but K is small, or you can't store everything.

Template#

Here is the weakest-link rule as code: a size-k min-heap that keeps the K largest items by any key. Flip the sign of the key to get the K smallest.

heap = empty min-heap of (key, index, item)   # heap[0] is the weakest
for each item, with its index:
    entry = (key of item, index, item)        # index breaks ties
    if len(heap) < k:
        push entry onto heap
    else if entry > heap[0]:
        replace heap[0] with entry            # evict the weakest: O(log k)
return the items in heap                      # any order

After the loop, heap[0] is the Kth largest. When the key is a distance, compare squared distances: if $A < B$ then $A^2 < B^2$, so $x^2 + y^2$ gives the same order with no square root, float cost or precision issue.

Other shapes#

Two more shapes solve some Top K problems faster or with ordered output.

Heapify, then pop K#

When all the data is already in memory, heapify the whole list ($O(N)$) and pop K times ($O(K \log N)$). The total, $O(N + K \log N)$, is much better than pushing items one at a time ($O(N \log N)$). It is also the easy way to get ordered output, because pops come out in order and tuple comparison breaks ties.

Bucket sort for frequencies#

For "top K by frequency", counts can't exceed N, so you can skip the heap. Put each value in the bucket for its count, then walk the buckets from the highest count down.

count each value
buckets = list of N + 1 empty lists          # index = frequency
for each (value, freq) in counts:
    append value to buckets[freq]
result = empty list
for freq from N down to 1:
    for each value in buckets[freq]:
        append value to result
        if len(result) == k: return result

The nested loop looks like $O(N^2)$, but it isn't: each unique value sits in exactly one bucket and is visited once, so the walk is $O(N)$. It does not give alphabetical order within a bucket.

Counting sort for Kth largest#

If values sit in a small range, count them in an array indexed by num - min_value and walk down from the maximum until K items are used up. That's $O(N + \text{range})$ with no heap at all.

Which heap pattern?#

Most heap problems fall into one of these shapes.

A size-k heap (Top K) and a k-way merge are different patterns: one caps the heap at K items, the other keeps one item per source.

Pitfalls#

These are the slips that come up most often with heaps.

• Heapifying everything and calling it "size K". It's correct, but costs $O(N)$ memory. When the heap starts from initial data, trim it to K there too, not only on later inserts.
• Forgetting to negate back when you read from a negated max-heap.
• Tuples with non-comparable items. (priority, obj) raises TypeError on a tie when obj has no ordering (a ListNode, a dict). Add an index as a tie-breaker, or wrap the object in a class with __lt__.
• Picking the wrong heap for Top K. K largest needs a min-heap, not a max-heap.
• Re-sorting in a loop. Sorting again every round of a simulation costs $O(n^2 \log n)$; a heap makes it $O(n \log n)$.

Problems in this pattern#

These problems use this pattern. Solve them before you open the hints below. Start with the first three.

• 1046. Last Stone Weight
• 703. Kth Largest Element in a Stream
• 215. Kth Largest Element in an Array
• 347. Top K Frequent Elements
• 973. K Closest Points to Origin
• 451. Sort Characters By Frequency
• 692. Top K Frequent Words

Hints#

Open a hint only after you have tried the problem for a while.

Key takeaways#

• heapq is a min-heap; negate values for a max-heap, and negate back when you read.
• K largest uses a min-heap of size K (the top is the weakest link), giving $O(N \log K)$.
• heapify is $O(N)$, so heapify-then-pop-K is $O(N + K \log N)$.
• Frequency problems can use bucket sort in $O(N)$ unless ties need ordering.
• Put an index in the tuple (or define __lt__) whenever the payload can't be compared.