Heaps and priority queues in Rust

Picture a hospital emergency room. You arrive with a sprained ankle, and you wait. Ten minutes later someone arrives with chest pain. They are seen first, even though you were there before them.

That is not a normal queue. A normal queue (lesson 5) is first come, first served. This one is most urgent first.

Computer scientists call it a priority queue. Every item carries a priority, which is a number that says how urgent it is. You can add an item at any time. When you ask for the next item, you always get the one with the highest priority.

A heap is the structure that makes a priority queue fast. Reach for one whenever you keep asking "what is the best item so far?": the next task to run, the nearest place to visit, the five biggest scores in a stream of millions.

A heap is a tree with one rule. In a min-heap, every parent is no bigger than its children. (A max-heap flips it: every parent is no smaller.) Follow the rule all the way up, and the smallest value must be sitting at the very top.

Notice what the rule does not say. It does not sort the whole tree. Siblings can be in any order, and 40 sits above 25 in the picture. The heap only keeps promise one: the best item is on top. That is a much weaker promise than "everything is sorted", and weaker promises are cheaper to keep.

The tree is also complete, meaning every level is full, except maybe the last, which fills from the left. That shape lets a heap live in a plain Vec, with no pointers at all. Read the tree level by level, left to right, and you get the row in the picture. Two small formulas find family members from a slot number i:

• the parent of slot i is at (i - 1) / 2
• the children of slot i are at 2i + 1 and 2i + 2

Check it on slot 2, which holds 15. Its children should be at 2*2 + 1 = 5 and 2*2 + 2 = 6. Those hold 25 and 50, exactly as drawn.

Now the two moves. Start with a tiny example, then the rule.

Push 5. Put it in the next free slot, at the end of the row. That is slot 7, under 40. But 5 is smaller than 40, so it swaps places with it. Then it is smaller than 20, so it swaps again. Then it is smaller than 10, and swaps once more. Now it sits on top. It floated up three levels.

Pop the top. You cannot just remove slot 0, because that leaves a hole. So take the last value in the row and drop it into the hole. It is probably too big for the top, so it sinks. At each step it swaps with the smaller of its two children, until both children are no smaller than it.

Every swap moves one level. A tree with n values is only about log n levels tall, so both moves take O(log n) steps. That is a handful of steps even for a million values.

Remember: a heap does not keep everything sorted. It keeps the best item on top, and fixes that cheaply after every change.

This heap holds i32 values and keeps the smallest on top.

/// A min-heap of `i32`: the smallest value is always at index 0.
/// The values live in a `Vec`, laid out as a tree level by level.
pub struct MinHeap {
    items: Vec<i32>,
}

impl MinHeap {
    pub fn new() -> Self {
        MinHeap { items: Vec::new() }
    }

    pub fn len(&self) -> usize {
        self.items.len()
    }

    pub fn is_empty(&self) -> bool {
        self.items.is_empty()
    }

    /// The smallest value, without removing it.
    pub fn peek(&self) -> Option<i32> {
        self.items.first().copied()
    }

    /// Adds a value at the end, then lets it float up to where it belongs.
    pub fn push(&mut self, value: i32) {
        self.items.push(value);
        let mut i = self.items.len() - 1;
        while i > 0 {
            let parent = (i - 1) / 2;
            if self.items[i] >= self.items[parent] {
                break; // it is no smaller than its parent, so it can stay
            }
            self.items.swap(i, parent);
            i = parent;
        }
    }

    /// Removes the smallest value. The last value takes its place and sinks.
    pub fn pop(&mut self) -> Option<i32> {
        if self.items.is_empty() {
            return None;
        }
        // Swapping first means `pop` on the Vec is the cheap end-removal.
        let last = self.items.len() - 1;
        self.items.swap(0, last);
        let smallest = self.items.pop();
        self.sink(0);
        smallest
    }

    /// Moves the value at `i` down until both children are no smaller.
    fn sink(&mut self, mut i: usize) {
        let n = self.items.len();
        loop {
            let left = 2 * i + 1;
            let right = left + 1;
            let mut smallest = i;
            if left < n && self.items[left] < self.items[smallest] {
                smallest = left;
            }
            if right < n && self.items[right] < self.items[smallest] {
                smallest = right;
            }
            if smallest == i {
                return;
            }
            self.items.swap(i, smallest);
            i = smallest;
        }
    }
}

impl Default for MinHeap {
    fn default() -> Self {
        Self::new()
    }
}

Go through it in order.

• items is the Vec from the picture. There is no node type, because the slot numbers are the links.
• peek looks at slot 0 and copies it out. first() gives an Option, which is None for an empty heap.
• push adds to the end, then loops. While the new value is smaller than its parent, it swaps with it and moves up. The moment it is no smaller, it stops. The break is the whole reason pushing is quick.
• pop swaps the top with the last value, and then calls pop on the Vec, which removes the last slot cheaply. What is left at the top is the value that moved in, so sink fixes it.
• sink looks at slot i and its two children, finds the smallest of the three, and swaps if a child won. The checks left < n and right < n stop it reading past the end of the Vec.

Here it is in use.

let mut heap = MinHeap::new();
for value in [50, 20, 40, 10, 30] {
    heap.push(value);
}

println!("{:?}", heap.peek());
assert_eq!(heap.peek(), Some(10));
assert_eq!(heap.pop(), Some(10));
assert_eq!(heap.pop(), Some(20));
assert_eq!(heap.len(), 3);

The values went in as 50, 20, 40, 10, 30 and came out smallest first. The Playground link on the code above runs this.

The standard library has BinaryHeap. It is the same idea, finished and well tested. One surprise: it is a max-heap, so pop gives the biggest item. That suits our hospital, where a bigger number means more urgent.

use std::collections::BinaryHeap;

// The biggest number is the most urgent, and it comes out first.
let mut waiting = BinaryHeap::new();
waiting.push((2, "sprained ankle"));
waiting.push((9, "chest pain"));
waiting.push((5, "broken arm"));

assert_eq!(waiting.peek(), Some(&(9, "chest pain")));
let (urgency, patient) = waiting.pop().unwrap();
println!("{urgency}: {patient}");
assert_eq!(waiting.pop(), Some((5, "broken arm")));
assert_eq!(waiting.len(), 1);

Each item here is a tuple of an urgency and a name. Rust compares tuples from left to right, so the urgency decides, and the name only matters if two urgencies are equal. peek borrows the top without removing it, so it returns Option<&T>, and pop takes it out and returns Option<T>.

What if you want the smallest first? There is no MinHeap type. Instead you wrap each item in Reverse, a small standard-library type that flips the comparison. A bigger number now counts as "smaller", so it sinks, and the true smallest rises.

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

let mut heap = BinaryHeap::new();
for n in [30, 10, 50, 20] {
    heap.push(Reverse(n)); // Reverse flips the order, so the smallest wins
}

let Reverse(smallest) = heap.pop().unwrap(); // unwrap the Reverse again
println!("{smallest}");
assert_eq!(smallest, 10);
assert_eq!(heap.peek(), Some(&Reverse(20)));

Reverse(n) wraps the number, so what comes back from pop is wrapped too. The line let Reverse(smallest) = ... unwraps it again. It is a tuple struct (a struct whose fields have no names), so .0 works too.

Remember: BinaryHeap gives the biggest first. For smallest first, wrap each item in Reverse.

Two more facts worth knowing. An item must have a total order, meaning any two items can be compared. i32 and String can. f64 cannot, because of NaN ("not a number"), which is not smaller or bigger than anything. And building a heap from a whole Vec in one go is cheaper than pushing one value at a time:

use std::collections::BinaryHeap;

// Building from a Vec in one go is cheaper than pushing one by one.
let heap = BinaryHeap::from(vec![4, 1, 7, 3]);
assert_eq!(heap.peek(), Some(&7));
println!("{:?}", heap.into_sorted_vec());

into_sorted_vec hands back everything in ascending order. Do not loop over a heap with iter() and expect order. It visits the values in storage order, which is only partly sorted.

The row to remember is the last-but-one. A heap is excellent at "give me the best" and poor at "is this value in here?". If you need to search, reach for a hash map or a sorted structure.

Each starter has a function signature and a todo!(), which is a placeholder that panics if it is ever run. Replace it with your code. The Playground link on each starter runs a small demo, so you can try your answer straight away.

pub fn last_stone_weight(stones: Vec<i32>) -> i32 {
    todo!("smash the two heaviest stones, over and over")
}

pub fn last_stone_weight(stones: Vec<i32>) -> i32 {
    use std::collections::BinaryHeap;

    // A max-heap hands over the heaviest stone each time.
    let mut heap = BinaryHeap::from(stones);
    while heap.len() > 1 {
        let (Some(first), Some(second)) = (heap.pop(), heap.pop()) else {
            break;
        };
        if first != second {
            heap.push(first - second); // the leftover chunk goes back in
        }
    }
    heap.pop().unwrap_or(0)
}

pub fn kth_largest(nums: &[i32], k: usize) -> Option<i32> {
    todo!("keep a min-heap of the k biggest values seen so far")
}

pub fn kth_largest(nums: &[i32], k: usize) -> Option<i32> {
    use std::cmp::Reverse;
    use std::collections::BinaryHeap;

    if k == 0 || k > nums.len() {
        return None;
    }
    // A min-heap capped at k values: its smallest is the k-th largest overall.
    let mut heap = BinaryHeap::new();
    for &n in nums {
        heap.push(Reverse(n));
        if heap.len() > k {
            heap.pop(); // drop the smallest, which cannot be in the top k
        }
    }
    heap.peek().map(|&Reverse(n)| n)
}

pub fn merge_sorted(lists: &[Vec<i32>]) -> Vec<i32> {
    todo!("always take the smallest front value from any list")
}

pub fn merge_sorted(lists: &[Vec<i32>]) -> Vec<i32> {
    use std::cmp::Reverse;
    use std::collections::BinaryHeap;

    // Each entry is (value, which list, position in that list).
    let mut heap = BinaryHeap::new();
    for (list_id, list) in lists.iter().enumerate() {
        if let Some(&first) = list.first() {
            heap.push(Reverse((first, list_id, 0)));
        }
    }

    let mut merged = Vec::new();
    while let Some(Reverse((value, list_id, pos))) = heap.pop() {
        merged.push(value);
        // Replace what we took with the next value from the same list.
        if let Some(&next) = lists[list_id].get(pos + 1) {
            heap.push(Reverse((next, list_id, pos + 1)));
        }
    }
    merged
}

Heap questions have a recognisable smell. A few habits help.

• Spot the signals. "The k largest", "the k smallest", "the next best", "merge k sorted things", and "keep a running median" all point to a heap.
• Know the size-k trick. To find the k largest, keep a min-heap of size k, as in problem 2. It feels backwards, but the smallest of the winners is exactly the one that gets thrown out. It costs O(n log k), not O(n log n).
• Say which heap you need. Rust's BinaryHeap is a max-heap. For smallest-first, say out loud that you will wrap items in Reverse. It shows you know the standard library.
• Be honest about searching. A heap cannot find an arbitrary value quickly, and it has no cheap way to change an item's priority. If a problem needs either, mention a different structure, or the trick of pushing a new entry and ignoring the stale one when it appears.
• Mention where heaps are used. The shortest-path algorithm called Dijkstra's, task schedulers, and event simulations all lean on a priority queue. You do not need to code them. Knowing the connection makes your answer sound less like a trick.
• Say the cost. "Push and pop are O(log n), peek is O(1), and building from n values is O(n)" is a complete answer.

Made by VeraStack Labs