Heap Patterns

Top K Pattern

Find the $K^{th}$ largest (or smallest) element in an array or stream. Alternatively: Find the “Top K” frequent elements, “Top K” closest points, etc.

We avoid sorting the entire array ($O(N\log N)$). Instead, we use a Heap to organize the data based on priority.

There are two main strategies:

Min-Heap (Best for “Largest” items)

• Concept: Maintain a heap of size exactly K.
• Why Min-Heap?: We want the K largest numbers. To do that, we must discard the small numbers. A Min-Heap gives us easy access to the smallest number (the root) so we can kick it out.
• Result: The root is the “Smallest of the Top K”, which is the $K^{th}$ largest overall.
• Complexity: $O(N\log K)$.

Max-Heap (Best for “Smallest” items or Static Arrays)

• Concept: Dump everything into a Max-Heap. Pop $K-1$ times.
• Why Max-Heap?: The root is the largest. If we remove the largest, then the 2nd largest… we eventually expose the $K^{th}$ largest.
• Complexity: $O(N+K\log N)$ (if using efficient Heapify).

Implementation

Rust defaults to Max-Heap and python defaults to Min-Heap

import heapq
 
# Python using Strategy A
class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        # 1. Heapify the list in O(N)
        heapq.heapify(nums)
 
        # 2. Pop the smallest elements until we only have K left.
        # The remaining K elements are the largest K in the array.
        while len(nums) > k:
            heapq.heappop(nums)
 
        # 3. The root is the smallest of the "Top K", i.e., the Kth Largest.
        return nums[0]

use std::collections::BinaryHeap;
 
// Rust using Strategy B
impl Solution {
    pub fn find_kth_largest(nums: Vec<i32>, mut k: i32) -> i32 {
        // 1. Build the heap using 'from'.
        // Optimization: This runs in O(N), not O(N log N).
        let mut heap = BinaryHeap::from(nums);
 
        // 2. Pop k-1 times to remove the largest, 2nd largest, etc.
        while k > 1 {
            heap.pop();
            k -= 1;
        }
 
        // 3. The current top is the Kth largest.
        *heap.peek().unwrap()
    }
}

Complexity Analysis

Operation	Strategy A (Size K)	Strategy B (Size N)	Sorting
Time	$O(N\log K)$	$O(N+K\log N)$	$O(N\log N)$
Space	$O(K)$	$O(N)$	$O(1)$ or $O(N)$

When to use which?

1. Use Strategy A (Min-Heap Size K) when $K$ is small compared to $N$, or when handling a Stream (infinite data).
2. Use Strategy B (Max-Heap Size N) when you have a static array and efficient heapify (like Rust’s from) is available.

Merge K Sorted Lists

You are given an array of k linked-lists. Each linked-list is already sorted in ascending order. Merge all the linked-lists into one sorted linked-list and return it.

Example: Input:

[
  1->4->5,
  1->3->4,
  2->6
]

Output: 1->1->2->3->4->4->5->6

Method 1: Brute Force

The “I don’t care that inputs are sorted” approach.

1. Iterate through every node in every list.
2. Dump all values into a standard array (Python List / Rust Vec).
3. Sort the array (Time: $O(N\log N)$).
4. Create a new Linked List from scratch using the sorted values.

Critique

• Time: $O(N\log N)$. This is slower than optimal because we re-sort data that was already partially sorted.
• Space: $O(N)$. We must store copies of every value.
• Verdict: Easy to code, but fails the “Engineering” interview because it ignores the input properties.

Implementation

def mergeKLists_brute(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
    nodes = []
    # 1. Collect all values
    for l in lists:
        while l:
            nodes.append(l.val)
            l = l.next
    
    # 2. Sort (Expensive)
    nodes.sort()
    
    # 3. Rebuild
    dummy = ListNode(0)
    curr = dummy
    for val in nodes:
        curr.next = ListNode(val)
        curr = curr.next
    return dummy.next

Method 2: Min-Heap

The Optimal Approach.

Imagine 5 sorted lines of people. To find the overall shortest person, you only need to look at the front of each line. You don’t need to look at the back. The “Global Minimum” is guaranteed to be one of the K heads.

1. Init: Push the Head of every non-empty list into a Min-Heap.
   • Heap Size is at most K.
2. Loop:
   • Pop the smallest node from the heap.
   • If that node has a .next, push that Next Node into the heap.
3. End: When heap is empty, we are done.

Complexity Analysis

• Time: $O(N\log K)$
  • We iterate through all $N$ nodes.
  • For every node, we perform a heap push/pop.
  • The heap size is never larger than $K$. So each operation is $\log K$.
• Space: $O(K)$
  • The heap only holds the current “frontier” (one node from each list).

Implementation

Python: The “Tuple Comparison” Trap

In Python, if you put tuples in a heap, it compares element by element: (val, node). If two nodes have the same val (e.g., two lists start with 1), Python tries to compare node1 < node2.

Constraint: ListNode objects cannot be compared. This throws a TypeError.

Fix: Store (val, unique_id, node). The unique_id breaks ties so Python never compares the nodes.

import heapq
 
class Solution:
    def mergeKLists(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
        # Min-Heap stores: (value, unique_index, node)
        heap = []
        
        # Initialize Heap with heads
        for i, node in enumerate(lists):
            if node:
                # We use 'i' as a tie-breaker so ListNode is never compared
                heapq.heappush(heap, (node.val, i, node))
        
        dummy = ListNode(0)
        curr = dummy
        
        while heap:
            val, i, node = heapq.heappop(heap)
            curr.next = node
            curr = curr.next
            
            if node.next:
                heapq.heappush(heap, (node.next.val, i, node.next))
                
        return dummy.next

Rust: The “Ownership & Ord” Trap

Rust is strictly typed.

1. Ordering: ListNode does not implement Ord. We cannot put it directly in a BinaryHeap.
2. Ownership: If we put a Box<ListNode> into the heap, the heap owns it. It becomes hard to link it to the result list.

The Strategy:

Instead of putting the Node in the heap, we put the Index of the list in the heap.

• Heap stores: (Value, List_Index).
• We use a Min-Heap (via Reverse).
• When we pop (val, i), we know the smallest element is at lists[i]. We take it, and advance lists[i] to lists[i].next.

use std::collections::BinaryHeap;
use std::cmp::Reverse;
 
impl Solution {
    pub fn merge_k_lists(lists: Vec<Option<Box<ListNode>>>) -> Option<Box<ListNode>> {
        // We need a mutable version of lists to advance pointers
        let mut lists = lists; 
        
        // Heap stores (Value, Index in 'lists' vector)
        // Wrapped in Reverse to behave like Min-Heap
        let mut heap = BinaryHeap::new();
 
        for (i, list) in lists.iter().enumerate() {
            if let Some(node) = list {
                heap.push(Reverse((node.val, i)));
            }
        }
 
        let mut dummy_head = Box::new(ListNode::new(0));
        let mut curr = &mut dummy_head;
 
        while let Some(Reverse((val, i))) = heap.pop() {
            // Logic:
            // a. We know lists[i] has the smallest value.
            // b. We take that node out of the vector (using .take()).
            // c. We make curr.next point to it.
            // d. If that node had a .next, we put it back into lists[i] and push to heap.
            
            if let Some(mut node) = lists[i].take() {
                let next_node = node.next.take();
                lists[i] = next_node;
                
                // If there is a next node, push its value and index back to heap
                if let Some(ref next) = lists[i] {
                    heap.push(Reverse((next.val, i)));
                }
                
                curr.next = Some(node);
                curr = curr.next.as_mut().unwrap();
            }
        }
 
        dummy_head.next
    }
}

(Note: The Rust implementation for Linked Lists is significantly more verbose due to strict memory safety. The logic above avoids implementing custom structs by leveraging the lists vector indices