Heap-Based Algorithms

A BinaryHeap (priority queue) dynamically maintains the minimum or maximum of a changing set of elements in $O(\log n)$ time per update, powering classic algorithms such as K-way sorted stream merging, online running medians (Two-Heap pattern), and interval scheduling.

Two-Heap Running Median Pattern:

  Smaller Half (`max_heap: BinaryHeap<i64>`)    Larger Half (`min_heap: BinaryHeap<Reverse<i64>>`)
  [ ... 10, 20, 30 ] <--- max is 30             min is 40 ---> [ 40, 50, 60 ... ]
                 \                             /
                  +--- Median sits right at --+
                       the two heap roots in O(1)!

1. Online Running Median with Two Heaps ($O(\log n)$ Insert, $O(1)$ Query)

To maintain the median of a dynamically growing stream of numbers:

• Store the smaller half of the numbers in a Max-Heap (BinaryHeap<i64>), so its .peek() gives the largest element of the lower half.
• Store the larger half of the numbers in a Min-Heap (BinaryHeap<Reverse<i64>>), so its .peek() gives the smallest element of the upper half.
• Balance the sizes so max_heap.len() is either equal to min_heap.len() or 1 element larger.

use std::cmp::Reverse;
use std::collections::BinaryHeap;

#[derive(Default)]
pub struct RunningMedian {
    lower_max: BinaryHeap<i64>,          // Max-heap of the smaller half
    upper_min: BinaryHeap<Reverse<i64>>, // Min-heap of the larger half
}

impl RunningMedian {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn push(&mut self, num: i64) {
        // 1. Push to lower half first, then move largest of lower half to upper half:
        self.lower_max.push(num);
        if let Some(largest_lower) = self.lower_max.pop() {
            self.upper_min.push(Reverse(largest_lower));
        }

        // 2. Keep `lower_max.len() >= upper_min.len()`:
        if self.upper_min.len() > self.lower_max.len() {
            if let Some(Reverse(smallest_upper)) = self.upper_min.pop() {
                self.lower_max.push(smallest_upper);
            }
        }
    }

    pub fn median_f64(&self) -> Option<f64> {
        let low_top = *self.lower_max.peek()? as f64;
        if self.lower_max.len() > self.upper_min.len() {
            Some(low_top)
        } else {
            let high_top = self.upper_min.peek()?.0 as f64;
            Some((low_top + high_top) * 0.5)
        }
    }
}

2. Merging $K$ Sorted Slices in $O(N \log K)$ Time

Given $K$ sorted slices containing $N$ total elements, a min-heap holding at most $K$ items Reverse((value, list_idx, elem_idx)) merges all $K$ lists into a single sorted vector in $O(N \log K)$ time and $O(K)$ auxiliary heap space:

pub fn merge_k_sorted(lists: &[&[i32]]) -> Vec<i32> {
    let total_len: usize = lists.iter().map(|s| s.len()).sum();
    let mut out = Vec::with_capacity(total_len);
    let mut min_heap: BinaryHeap<Reverse<(i32, usize, usize)>> =
        BinaryHeap::with_capacity(lists.len());

    // Seed heap with the first element of each non-empty list
    for (list_idx, &slice) in lists.iter().enumerate() {
        if let Some(&first_val) = slice.first() {
            min_heap.push(Reverse((first_val, list_idx, 0)));
        }
    }

    while let Some(Reverse((val, list_idx, elem_idx))) = min_heap.pop() {
        out.push(val);
        let next_idx = elem_idx + 1;
        if let Some(&next_val) = lists[list_idx].get(next_idx) {
            min_heap.push(Reverse((next_val, list_idx, next_idx)));
        }
    }

    out
}

fn main() {
    let mut tracker = RunningMedian::new();
    for x in [10, 40, 20] {
        tracker.push(x);
    }
    assert_eq!(tracker.median_f64(), Some(20.0));
    tracker.push(30);
    assert_eq!(tracker.median_f64(), Some(25.0));

    let merged = merge_k_sorted(&[&[1, 4, 7], &[2, 5, 8], &[3, 6, 9]]);
    assert_eq!(merged, vec![1, 2, 3, 4, 5, 6, 7, 8, 9]);

    println!("Running median: {:?}, K-way merged: {merged:?}", tracker.median_f64());
}

Running median: Some(25.0), K-way merged: [1, 2, 3, 4, 5, 6, 7, 8, 9]