38 / 387 min read

Heap Patterns

Priority queues for top-K, streaming median, and scheduling problems. Two heaps together give O(log n) median maintenance; a single heap gives O(n log k) top-K selection.

The Pattern

A heap (priority queue) answers the question "what is the current minimum/maximum?" in O(1) and updates in O(log n). Three canonical heap patterns appear across most interview problems:

  • Top-K: Maintain a min-heap of size k — the smallest element in the heap is always the kth largest overall.
  • Two heaps (median): A max-heap for the lower half, a min-heap for the upper half. Rebalance after each insertion so their sizes differ by at most one.
  • Scheduling / task ordering: Use a heap to greedily pick the best available task at each time step.

Rust's BinaryHeap is a max-heap. For min-heap behavior, wrap values in std::cmp::Reverse.

Kth Largest Element in a Stream

Maintain a min-heap of size k. The root is always the kth largest element seen so far. When a new element is added, push it and pop if size exceeds k.

use std::collections::BinaryHeap;
use std::cmp::Reverse;
 
struct KthLargest {
    k: usize,
    heap: BinaryHeap<Reverse<i32>>, // min-heap via Reverse
}
 
impl KthLargest {
    fn new(k: usize, nums: &[i32]) -> Self {
        let mut kl = KthLargest { k, heap: BinaryHeap::new() };
        for &n in nums { kl.add(n); }
        kl
    }
 
    fn add(&mut self, val: i32) -> i32 {
        self.heap.push(Reverse(val));
        if self.heap.len() > self.k { self.heap.pop(); }
        self.heap.peek().unwrap().0
    }
}
 
fn main() {
    let mut kl = KthLargest::new(3, &[4, 5, 8, 2]);
    println!("{}", kl.add(3));  // 4  (heap: [4,5,8])
    println!("{}", kl.add(5));  // 5  (heap: [5,5,8])
    println!("{}", kl.add(10)); // 5  (heap: [5,8,10])
    println!("{}", kl.add(9));  // 8  (heap: [8,9,10])
    println!("{}", kl.add(4));  // 8  (heap: [8,9,10])
}

K Closest Points to Origin

Find the k points closest to the origin. Use a max-heap of size k keyed by squared distance — if a point is closer than the furthest current top-k point, evict the farthest and add the new point.

fn k_closest(points: &[[i32; 2]], k: usize) -> Vec<[i32; 2]> {
    let mut heap: BinaryHeap<(i32, [i32; 2])> = BinaryHeap::new(); // max-heap by distance
 
    for &pt in points {
        let dist = pt[0] * pt[0] + pt[1] * pt[1];
        heap.push((dist, pt));
        if heap.len() > k { heap.pop(); }
    }
 
    heap.into_iter().map(|(_, pt)| pt).collect()
}
 
fn main() {
    let points = [[1,3],[-2,2],[5,8],[0,1]];
    let mut result = k_closest(&points, 2);
    result.sort();
    println!("{:?}", result); // [[-2,2],[0,1]]
}

Find K Pairs with Smallest Sums

Given two sorted arrays, find the k pairs (one element from each) with the smallest sums. Use a min-heap seeded with (nums1[i], nums2[0]) for all i. Each time you pop a pair (i, j), push (i, j+1) — advancing along nums2 for that fixed nums1[i] row.

fn k_smallest_pairs(nums1: &[i32], nums2: &[i32], k: usize) -> Vec<[i32; 2]> {
    if nums1.is_empty() || nums2.is_empty() { return vec![]; }
    let mut heap: BinaryHeap<Reverse<(i32, usize, usize)>> = BinaryHeap::new();
 
    for i in 0..nums1.len().min(k) {
        heap.push(Reverse((nums1[i] + nums2[0], i, 0)));
    }
 
    let mut result = Vec::new();
    while result.len() < k {
        let Some(Reverse((_, i, j))) = heap.pop() else { break };
        result.push([nums1[i], nums2[j]]);
        if j + 1 < nums2.len() {
            heap.push(Reverse((nums1[i] + nums2[j+1], i, j+1)));
        }
    }
    result
}
 
fn main() {
    println!("{:?}", k_smallest_pairs(&[1,7,11], &[2,4,6], 3));
    // [[1,2],[1,4],[1,6]]
 
    println!("{:?}", k_smallest_pairs(&[1,1,2], &[1,2,3], 2));
    // [[1,1],[1,1]]
}

Find Median from Data Stream — Two Heaps

Maintain a max-heap for the lower half and a min-heap for the upper half. After each insertion, rebalance so |lower| - |upper| <= 1. The median is either the max of lower (odd total) or the average of both tops (even total).

struct MedianFinder {
    lower: BinaryHeap<i32>,          // max-heap: lower half
    upper: BinaryHeap<Reverse<i32>>, // min-heap: upper half
}
 
impl MedianFinder {
    fn new() -> Self {
        MedianFinder { lower: BinaryHeap::new(), upper: BinaryHeap::new() }
    }
 
    fn add_num(&mut self, num: i32) {
        self.lower.push(num);
        // Balance: move max of lower to upper
        self.upper.push(Reverse(self.lower.pop().unwrap()));
        // Keep lower.len() >= upper.len()
        if self.lower.len() < self.upper.len() {
            self.lower.push(self.upper.pop().unwrap().0);
        }
    }
 
    fn find_median(&self) -> f64 {
        if self.lower.len() > self.upper.len() {
            *self.lower.peek().unwrap() as f64
        } else {
            let lo = *self.lower.peek().unwrap() as f64;
            let hi = self.upper.peek().unwrap().0 as f64;
            (lo + hi) / 2.0
        }
    }
}
 
fn main() {
    let mut mf = MedianFinder::new();
    mf.add_num(1);
    mf.add_num(2);
    println!("{}", mf.find_median()); // 1.5
    mf.add_num(3);
    println!("{}", mf.find_median()); // 2.0
    mf.add_num(4);
    println!("{}", mf.find_median()); // 2.5
}

Reorganize String

Rearrange a string so no two adjacent characters are the same. Greedily always place the most frequent remaining character. Use a max-heap keyed by frequency. After placing a character, hold it back for one round (can't place the same character consecutively), then re-add it.

fn reorganize_string(s: &str) -> String {
    let mut freq = [0i32; 26];
    for b in s.bytes() { freq[(b - b'a') as usize] += 1; }
 
    let mut heap: BinaryHeap<(i32, u8)> = freq.iter().enumerate()
        .filter(|&(_, &f)| f > 0)
        .map(|(i, &f)| (f, b'a' + i as u8))
        .collect();
 
    let mut result = Vec::new();
    while let Some((count, ch)) = heap.pop() {
        if result.last() == Some(&ch) {
            // Can't place ch — take the next best
            if let Some((count2, ch2)) = heap.pop() {
                result.push(ch2);
                if count2 - 1 > 0 { heap.push((count2 - 1, ch2)); }
                heap.push((count, ch)); // put ch back
            } else {
                return String::new(); // impossible
            }
        } else {
            result.push(ch);
            if count - 1 > 0 { heap.push((count - 1, ch)); }
        }
    }
    String::from_utf8(result).unwrap()
}
 
fn main() {
    println!("{}", reorganize_string("aab"));  // "aba"
    println!("{}", reorganize_string("aaab")); // "" (impossible)
    println!("{}", reorganize_string("vvvlo")); // "vlvov" or similar valid arrangement
}

Task Scheduler

Given a list of tasks (letters) and cooldown n (same task must be at least n apart), find the minimum total time to execute all tasks. Greedily at each time slot, run the most frequent available task. Track the next-available time for each task.

fn least_interval(tasks: &[char], n: i32) -> i32 {
    let mut freq = [0i32; 26];
    for &t in tasks { freq[(t as u8 - b'A') as usize] += 1; }
 
    // (freq, next_available_time)
    let mut heap: BinaryHeap<(i32, i32)> = freq.iter()
        .filter(|&&f| f > 0)
        .map(|&f| (f, 0i32))
        .collect();
 
    let mut time = 0;
    while !heap.is_empty() {
        // Collect tasks whose cooldown has expired
        let mut waiting = Vec::new();
        while let Some(&(f, avail)) = heap.peek() {
            if avail <= time { heap.pop(); break; }
            heap.pop();
            waiting.push((f, avail));
        }
        for item in waiting { heap.push(item); }
 
        if let Some((f, _)) = heap.pop() {
            if f - 1 > 0 { heap.push((f - 1, time + n + 1)); }
        }
        time += 1;
    }
    time
}
 
fn main() {
    let tasks: Vec<char> = "AABABC".chars().collect();
    println!("{}", least_interval(&tasks, 2)); // 8: A->B->C->A->B->idle->A->B
    let tasks2: Vec<char> = "AAABBB".chars().collect();
    println!("{}", least_interval(&tasks2, 2)); // 8: A->B->idle->A->B->idle->A->B
}

Sliding Window Maximum — Monotonic Deque (Heap Alternative)

Find the maximum in each window of size k. A max-heap works in O(n log n) but the optimal solution uses a monotonic deque in O(n). Shown here with a heap for contrast.

fn max_sliding_window(nums: &[i32], k: usize) -> Vec<i32> {
    // O(n log n) heap approach: lazy deletion via (value, index)
    let mut heap: BinaryHeap<(i32, usize)> = BinaryHeap::new();
    let mut result = Vec::new();
 
    for i in 0..nums.len() {
        heap.push((nums[i], i));
        // Lazily remove elements outside the window
        while heap.peek().map_or(false, |&(_, idx)| idx + k <= i) {
            heap.pop();
        }
        if i + 1 >= k {
            result.push(heap.peek().unwrap().0);
        }
    }
    result
}
 
fn main() {
    println!("{:?}", max_sliding_window(&[1,3,-1,-3,5,3,6,7], 3));
    // [3, 3, 5, 5, 6, 7]
    println!("{:?}", max_sliding_window(&[1], 1));
    // [1]
}