DEV Community

Timevolt
Timevolt

Posted on

Heaps & Priority Queues: The One Heap to Rule Them All

The Quest Begins (The “Why”)

I still remember the first time I faced a “Kth largest element” question in an interview. My brain went straight to the obvious: sort the array and pick the element at index len‑k. It worked, but the interviewer’s eyebrows rose when I said the runtime was O(n log n). “Can we do better?” they asked. I felt like a knight staring at a dragon’s scales, wondering if there was a hidden weak spot.

That moment kicked off a little obsession: what if we could keep track of only the top k items while scanning the array once? The answer lived in a data structure I’d always taken for granted—the heap.

The Revelation (The Insight)

Why a heap?

A heap is a binary tree where every parent is either ≤ (min‑heap) or ≥ (max‑heap) its children. The beautiful part is that the root always holds the extreme value (minimum or maximum) of the whole set. If we can get that extreme in O(1) time and update the structure in O(log n) time, we’ve got a super‑charged priority queue.

The “keep‑k‑largest” trick

Imagine we walk through the array and maintain a min‑heap of size k that contains the k largest numbers we’ve seen so far.

  • The heap’s root is the smallest of those k numbers.
  • When a new element x arrives:
    • If the heap has fewer than k items, we just push x.
    • If the heap already has k items and x > root, we pop the root (the current smallest of the top k) and push x.
    • Otherwise we ignore x.

At the end, the root of the heap is the kth largest element.

Why does this work?

Because we never discard a number that could belong to the final top k. Any element smaller than the current root is guaranteed to be smaller than at least k elements already in the heap, so it can’t possibly be the kth largest. Conversely, any element larger than the root must replace that root, preserving the invariant that the heap always holds the k largest seen so far.

Building a heap in O(n)

Most people think inserting n items one‑by‑one into a heap costs O(n log n). The surprise is that we can heapify an arbitrary array in linear time. The intuition: heapify works from the bottom up. Most nodes are leaves (height 0) and need no work; only a fraction of nodes have height 1, fewer have height 2, and so on. Summing the work over all levels yields a geometric series that collapses to O(n).

That’s why the heap feels like finding a hidden cheat code—it gives us both fast retrieval and cheap construction.

Wielding the Power (Code & Examples)

Before: the naïve sort

def kth_largest_sort(nums, k):
    nums.sort()               # O(n log n)
    return nums[-k]
Enter fullscreen mode Exit fullscreen mode

Simple, but the sort step dominates the runtime for large inputs.

After: min‑heap of size k

import heapq

def kth_largest_heap(nums, k):
    # Build a min‑heap with the first k elements
    min_heap = nums[:k]
    heapq.heapify(min_heap)          # O(k) → O(n) when k is proportional to n

    # Process the rest
    for num in nums[k:]:
        if num > min_heap[0]:        # only interesting if it beats the current kth
            heapq.heapreplace(min_heap, num)  # pop root & push num in O(log k)

    return min_heap[0]               # the root is the kth largest
Enter fullscreen mode Exit fullscreen mode

Why this is faster:

  • Heapifying the first k items is O(k).
  • Each of the remaining n‑k items triggers at most one heapreplace, which is O(log k).
  • Total time: O(k + (n‑k)·log k). When k is small relative to n (the common interview case), this is essentially O(n log k), a huge win over O(n log n). If k is a constant, it collapses to pure O(n).

Another classic: merging k sorted lists

Problem: Given k sorted arrays, produce one sorted array.

Naïve approach: concatenate and sort → O(N log N) where N is total elements.

Heap solution:

def merge_k_lists(lists):
    min_heap = []
    # Initialize heap with the first element of each list
    for i, lst in enumerate(lists):
        if lst:                                 # guard against empty lists
            heapq.heappush(min_heap, (lst[0], i, 0))  # (value, list_id, index_in_list)

    result = []
    while min_heap:
        val, lst_idx, elem_idx = heapq.heappop(min_heap)
        result.append(val)
        # Push the next element from the same list, if any
        if elem_idx + 1 < len(lists[lst_idx]):
            next_val = lists[lst_idx][elem_idx + 1]
            heapq.heappush(min_heap, (next_val, lst_idx, elem_idx + 1))

    return result
Enter fullscreen mode Exit fullscreen mode

Why it works: The heap always holds the smallest unseen element from each list, so popping gives the next global minimum in O(log k). Each of the N elements is pushed and popped once → O(N log k).

Traps to avoid (the “boss fights”)

Trap What happens Fix
Forgetting to heapify after slicing the first k items You treat a plain list as a heap → pop/push give wrong order Call heapq.heapify right after creating the heap
Using a max‑heap when you need a min‑heap (or vice‑versa) The root no longer represents the extreme you need Remember: Python’s heapq is a min‑heap; for a max‑heap push negative values or use -val
Not handling empty input or k > len(nums) Index errors or returning garbage Guard with if not nums or k > len(nums): raise ValueError
Pushing every element without the if num > heap[0] check Heap size grows beyond k → O(n log n) again Keep the size invariant strict

Why This New Power Matters

Once you internalize the heap‑as‑priority‑queue mindset, a whole class of problems stops feeling like grinding and starts feeling like crafting a spell:

  • Streaming analytics: keep the top‑k trends in real time with bounded memory.
  • Event simulation: process events in timestamp order using a priority queue (think of a game’s event loop).
  • Graph algorithms: Dijkstra’s and Prim’s become trivial when you reach for a heap.

The best part? The data structure is tiny—just an array—and the operations are battle‑tested in every language’s standard library. You’ll find yourself reaching for it instinctively, like a reflex, whenever you need to extract “the best so far.”

Your Turn

Here’s a little challenge to cement the power: Implement the sliding window maximum (given an array and window size w, return the max of each contiguous sub‑array of length w) using a max‑heap. Think about how you’d discard elements that fall out of the window efficiently.

Drop your solution in the comments, tweet it, or just try it in your REPL. I can’t wait to see how you wield the heap!

Happy coding, and may your roots always be strong. 🚀

Top comments (0)