Learn interview-prep-45

Heaps and Priority Queues

Browsing without an account — progress is saved in this browser only.Sign in to save it

Heaps and Priority Queues

Reading120 min

Think of a hospital emergency room. Patients don't get treated in the order they arrived; the most critical case always goes next, no matter where they are in the waiting line. A heap is that triage system for data: it always hands you the smallest (or largest) item first, cheaply, even as new items keep arriving. That's the engine behind top-K problems, merging sorted data, scheduling, and streaming statistics.

A heap's shape: a complete binary tree in an array

A heap is a complete binary tree stored in a plain array: every level is fully filled except possibly the last, which fills left to right with no gaps. That completeness is exactly what lets you store the tree in an array with no pointers wasted, since a node's children live at predictable index offsets.

For a node at index i (0-indexed array):

def parent(i):  return (i - 1) // 2
def left(i):    return 2 * i + 1
def right(i):   return 2 * i + 2

Array: [1, 3, 2, 7, 5, 4, 8]
Tree:
              1(0)
            /      \
        3(1)         2(2)
       /    \        /    \
    7(3)   5(4)   4(5)   8(6)

Because the tree is always complete, its height is always O(log n) for n elements. That height bound is exactly why every heap operation costs O(log n), not more.

Remember: a heap is a complete binary tree living inside a flat array. Index math (2i+1, 2i+2, (i-1)//2) replaces pointers.

Min heap vs max heap

Min heap: every parent is less than or equal to both its children. The smallest element always sits at the root, index 0. Python's heapq module implements a min heap.

Max heap: every parent is greater than or equal to both children. The largest element always sits at the root.

import heapq

min_heap = []
for x in [5, 1, 8, 3]:
    heapq.heappush(min_heap, x)
print(min_heap[0])  # 1 — smallest, always at index 0

# Python has no max-heap; negate values as the standard workaround
max_heap = []
for x in [5, 1, 8, 3]:
    heapq.heappush(max_heap, -x)
print(-max_heap[0])  # 8 — largest

Pitfall with the negation trick: when you push tuples for a max heap, such as (priority, item), negate only the priority field, and remember to negate it back when you pop. When priorities tie, Python compares the next tuple element. If that's an unorderable object, like a custom class, you get a TypeError, so add an index or unique tiebreaker as a third tuple element to sidestep this.

	Min heap	Max heap
Root	Smallest	Largest
`heapq` support	Native	Simulate via negation
Use for	Top-K largest (keep heap size K), Dijkstra, merge-K-lists	Top-K smallest, task scheduling by highest priority first

Remember: Python's heapq is min-heap only. For a max heap, negate values going in and negate the result coming out.

Building a heap from scratch in O(n)

heapify builds a valid heap from an arbitrary array in O(n) time, not the O(n log n) you might guess from "n inserts of O(log n) each." The trick: start from the last non-leaf node and sift each node down, working backward toward the root. Most nodes sit near the bottom of the tree, where a sift-down is cheap, so the total work sums to O(n), not O(n log n).

def sift_down(arr, i, n):
    while True:
        smallest = i
        l, r = 2 * i + 1, 2 * i + 2
        if l < n and arr[l] < arr[smallest]:
            smallest = l
        if r < n and arr[r] < arr[smallest]:
            smallest = r
        if smallest == i:
            break
        arr[i], arr[smallest] = arr[smallest], arr[i]
        i = smallest

def heapify(arr):
    n = len(arr)
    for i in range(n // 2 - 1, -1, -1):  # start at last non-leaf, go to root
        sift_down(arr, i, n)

data = [9, 4, 7, 1, 3, 8]
heapify(data)
print(data)  # a valid min-heap array, e.g. [1, 3, 7, 9, 4, 8]

Python's heapq.heapify(list) does exactly this in place, in O(n).

Remember: heapify is O(n), not O(n log n), because it starts from the bottom, where most nodes live and sifting is cheap.

Kth Largest Element in Array

LeetCode 215, Heap, QuickSelect alternative

Keep a min heap that never holds more than k elements. After scanning every number, the heap's root is the k-th largest: everything smaller has already been evicted, leaving exactly the top k values, the smallest of which sits at the root.

Approach: Push each number; if the heap grows past k, pop the smallest. At the end, the root is the answer.

def findKthLargest(nums: list[int], k: int) -> int:
    heap = []
    for num in nums:
        heapq.heappush(heap, num)
        if len(heap) > k:
            heapq.heappop(heap)
    return heap[0]

Complexity: Time O(n log k), space O(k), much better than sorting's O(n log n) when k is small. QuickSelect reaches average O(n) but worst-case O(n²) and mutates the input, so the heap approach is the safer one to lead with, mentioning QuickSelect as the follow-up optimization.

Common mistakes: Using a max heap of all n elements and popping k times: O(n + k log n), but it needlessly builds a heap over the whole array when a size-k min heap is smaller and easier to reason about. Forgetting the size check lets the heap grow unbounded.

Remember: a size-k min heap; whatever survives eviction, its root is the k-th largest.

Top K Frequent Elements

LeetCode 347, Heap, Frequency counting

Count frequencies first; then you need the k elements with the highest counts. This is the same "maintain a size-k heap" pattern as Kth Largest, just ordered by frequency instead of raw value.

Approach: Counter for frequencies, then a min heap of size k keyed by frequency, evicting the lowest-frequency entry whenever the heap exceeds k.

from collections import Counter

def topKFrequent(nums: list[int], k: int) -> list[int]:
    counts = Counter(nums)
    heap = []
    for num, freq in counts.items():
        heapq.heappush(heap, (freq, num))
        if len(heap) > k:
            heapq.heappop(heap)
    return [num for freq, num in heap]

Complexity: Time O(n log k), space O(n) for the counter. heapq.nlargest(k, counts.keys(), key=counts.get) is the idiomatic one-liner if the interviewer allows library shortcuts, but be ready to write it by hand, since that's usually the actual ask.

Common mistakes: Sorting the full frequency map, O(n log n), when k is small and a heap does better. Forgetting the answer only needs the numbers, not the frequencies, when building the return list.

Every "find the K elements with the highest count" problem builds on the same first step:

Python

from collections import Counter
freq = Counter(arr)  # O(n) — {element: count}

From there, three approaches, from simplest to fastest:

Approach	How	Time	When to use
A. Sort	`sorted(freq.items(), key=lambda x: -x[1])[:k]`	O(n log n)	Default, safe answer. Correct and easy to write under pressure.
B. Size-k heap	The pattern above: push, evict past k	O(n log k)	The standard "top-k" answer once asked to optimize. Best when k is much smaller than n.
C. Bucket sort	Index directly by frequency (below)	O(n)	When pushed further: "can you beat O(n log k)?"

Approach C's insight: a value's frequency can never exceed n (the array's length), so instead of comparing or sorting, index an array directly by frequency.

Python

def top_k_frequent_bucket(arr, k):
    freq = Counter(arr)
    buckets = [[] for _ in range(len(arr) + 1)]  # index = frequency
    for num, count in freq.items():
        buckets[count].append(num)

    result = []
    for count in range(len(buckets) - 1, 0, -1):  # walk from highest frequency down
        for num in buckets[count]:
            result.append(num)
            if len(result) == k:
                return result
    return result

Bucket sort trades a general-purpose comparison sort for an array indexed by the bounded range of possible frequencies (0..n), which gives true linear time. It's the answer to give when an interviewer pushes for something better than O(n log n).

Remember: sort is the safe default, a size-k heap is the expected optimization, bucket sort is the answer when asked to beat O(n log k).

Merge K Sorted Lists

LeetCode 23, Heap, Merge technique

At any point in the merge, the next smallest overall value is the smallest among the current heads of all k lists. A min heap holding one candidate per list gives you that smallest value in O(log k), instead of scanning all k heads in O(k) at every step.

Approach: Push the head of each list into a min heap keyed by value (with a tiebreaker index, since comparing ListNode objects directly would fail). Pop the minimum, append it to the result, and push its .next if one exists.

class ListNode:
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next

def mergeKLists(lists: list[ListNode]) -> ListNode:
    heap = []
    for i, node in enumerate(lists):
        if node:
            heapq.heappush(heap, (node.val, i, node))  # i breaks value ties

    dummy = ListNode()
    tail = dummy

    while heap:
        val, i, node = heapq.heappop(heap)
        tail.next = node
        tail = tail.next
        if node.next:
            heapq.heappush(heap, (node.next.val, i, node.next))

    return dummy.next

Complexity: Time O(N log k), where N is the total node count across all lists. Space O(k) for the heap.

Common mistakes: Pushing (val, node) without a tiebreaker; comparing ListNode objects when values tie raises a type error. Forgetting the dummy-head pattern makes assembling the result list awkward.

Remember: the heap only ever holds k candidates, one per list's current head. Popping the smallest and pushing its successor keeps that invariant.

Find Median from Data Stream

LeetCode 295, Heap, Two heaps pattern

Split the stream into two halves: a max heap holding the smaller half, a min heap holding the larger half. Keep the two halves balanced in size, never differing by more than 1. The median is then either the max heap's root, the min heap's root, or the average of both, with no full sort needed on every insert.

Approach: Add each new number to the appropriate heap, then rebalance so the sizes never differ by more than 1.

class MedianFinder:
    def __init__(self):
        self.small = []  # max heap (negated), holds the smaller half
        self.large = []  # min heap, holds the larger half

    def addNum(self, num: int) -> None:
        heapq.heappush(self.small, -num)
        # ensure every element in small <= every element in large
        heapq.heappush(self.large, -heapq.heappop(self.small))
        if len(self.large) > len(self.small):
            heapq.heappush(self.small, -heapq.heappop(self.large))

    def findMedian(self) -> float:
        if len(self.small) > len(self.large):
            return -self.small[0]
        return (-self.small[0] + self.large[0]) / 2.0

Complexity: Time O(log n) per addNum, O(1) per findMedian. Space O(n) to hold the stream.

Common mistakes: Forgetting to route every insert through small first. The push-then-pop-then-rebalance sequence above guarantees correctness even when the new number actually belongs in large. Allowing a size imbalance greater than 1, which breaks the median formula. A plain sorted list with insertion also works, with O(n) insert, but it doesn't scale as well as the two-heap approach for large streams; mention it as the naive baseline.

Remember: two heaps, split at the midpoint, rebalanced after every insert so sizes never differ by more than 1. The median is at the two roots.

Heaps and Priority Queues

A heap's shape: a complete binary tree in an array

Min heap vs max heap

Building a heap from scratch in O(n)

Kth Largest Element in Array

Top K Frequent Elements

The full menu: three ways to find top K, in order of when to reach for each

Merge K Sorted Lists

Find Median from Data Stream