Heaps

Chapter 3 used vectors, strings, and hash maps. This chapter adds one more collection, the heap. A heap answers one question quickly: “what is the largest item right now?” It keeps answering it while items are added and removed.

That question comes up whenever you want the best few items from many. Two examples are the most frequent words in a text and the next job to run from a queue ordered by priority. You will build three programs. Two find the most frequent items. The third keeps the median of a stream of numbers.

As in chapter 3, I start with how the structure works, with pictures, and then show the code.

A few terms used throughout the chapter:

• A binary tree is a set of nodes where each node has at most two children, a left one and a right one. The node with no parent is the root.
• A max-heap is a binary tree where every parent is at least as large as each of its children. So the root holds the largest item.
• A min-heap is the same with “at least as large” replaced by “at most as large”. The root holds the smallest item.
• The median of a list of numbers is the middle value once they are sorted. For an even count, it is the average of the two middle values.

5.1 A tree stored in a vector

A heap is a tree, but it does not store pointers between nodes. It stores the nodes in a Vec, level by level: the root first, then its two children, then their four children, and so on. Figure 5.1 shows a max-heap of six numbers and the vector that holds it.

Because the nodes are packed level by level, simple arithmetic finds a node’s relatives. The children of index i are at 2i + 1 and 2i + 2. Its parent is at (i − 1) / 2, rounded down. In figure 5.1, the children of 7 (index 1) are at indexes 3 and 4, which hold 3 and 5.

The heap rule only compares parents with children. It says nothing about brothers or cousins. In figure 5.1, 7 is on the left and 8 on the right, and that is allowed. That weak rule keeps a heap cheap to maintain.

A tree with n nodes packed this way has about log₂ n levels. A heap of a million items is only about 20 levels deep.

5.2 Adding and removing

5.2.1 Push

To add an item, put it at the end of the vector, which is the next free spot on the bottom level. It may be larger than its parent, which breaks the rule. So swap it with its parent, and keep swapping upward until it is not larger than its parent, or it reaches the root. Figure 5.2 pushes 10.

Each swap moves the item one level up, and there are about log₂ n levels. So a push takes O(log n) steps.

5.2.2 Pop

To remove the largest item, take the root. That leaves a hole at the top. Fill it with the last item in the vector, which shrinks the vector by one. The moved item is probably too small for the top. So swap it with its larger child, and keep swapping downward until both its children are smaller. Figure 5.3 pops the 10 that figure 5.2 added.

It swaps with the larger child for a reason. That child is about to become the parent of the other one, and a parent must be at least as large as its children. Pop is also O(log n).

Looking at the largest item without removing it is O(1): it is always at index 0.

5.3 Heaps in Rust

Rust’s standard library provides std::collections::BinaryHeap<T>, a max-heap. Its main methods match the operations above:

To decide which item is larger, BinaryHeap uses the Ord trait, which types such as i32, String, and tuples implement.

5.3.1 How tuples compare

Tuples compare element by element, from the left. Two tuples are compared on their first elements, and only if those are equal does the second element decide. So (3, 1) > (2, 9), because 3 > 2, and the 9 is never looked at. And (3, 1) < (3, 2), because the first elements tie and 1 < 2.

This gives a way to choose what a heap orders by: put that field first in a tuple. A heap of (count, value) pairs orders by count, and breaks ties by value.

5.3.2 A min-heap with Reverse

BinaryHeap has no min-heap mode. Instead, the standard library provides a wrapper, std::cmp::Reverse<T>, whose comparison is the opposite of T’s. Reverse(3) > Reverse(5), because 3 < 5. So a BinaryHeap<Reverse<i32>> keeps the smallest number on top.

To put a value in, wrap it: heap.push(Reverse(x)). To get it out, unwrap it with a pattern: if let Some(Reverse(x)) = heap.pop().

Now you can read the three programs.

5.4 The k most frequent numbers

The problem. Given a list of numbers and a count k, return the k numbers that occur most often. For [1, 1, 1, 2, 2, 3] and k = 2, the answer is 1 and 2.

The idea. Count each number with a HashMap. Put every (count, number) pair into a max-heap. Pop k times.

Listing 5.1 Top k with a full heap. src/problems/top_k_frequent.rs

//! Top K Frequent Elements.
//!
//! Count occurrences with a HashMap, then use a max-heap (`BinaryHeap`) to pull
//! out the top k elements. O(n log n) worst case; acceptable and readable.

use std::collections::{BinaryHeap, HashMap};

pub fn top_k_frequent(nums: &[i32], k: usize) -> Vec<i32> {
    let mut counts: HashMap<i32, usize> = HashMap::with_capacity(nums.len());
    for &num in nums.iter() {
        *counts.entry(num).or_insert(0) += 1;
    }

    // BinaryHeap is a max-heap in Rust, so (count, num) keeps highest counts first.
    let mut heap: BinaryHeap<(usize, i32)> = counts
        .into_iter()
        .map(|(num, count)| (count, num))
        .collect();

    let mut result = Vec::with_capacity(k);
    while result.len() < k {
        match heap.pop() {
            Some((_, num)) => result.push(num),
            None => break,
        }
    }
    result
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn returns_top_two() {
        let nums = [1, 1, 1, 2, 2, 3];
        let mut result = top_k_frequent(&nums, 2);
        result.sort_unstable();
        assert_eq!(result, vec![1, 2]);
    }

    #[test]
    fn k_larger_than_unique_values_returns_all() {
        let result = top_k_frequent(&[1, 2, 3], 10);
        assert_eq!(result.len(), 3);
    }
}

Step through it:

1. *counts.entry(num).or_insert(0) += 1; counts one occurrence. entry(num) finds the map slot for num, and or_insert(0) creates it with 0 if it is missing. Both steps return a mutable reference to the count, and * followed by += 1 adds one to the value it points to.
2. counts.into_iter().map(|(num, count)| (count, num)).collect() turns each map entry around, so the count comes first, and collects the pairs straight into a BinaryHeap. A BinaryHeap can be built with collect like a Vec.
3. The while loop pops until it has k numbers. If there are fewer than k distinct numbers, pop returns None and the loop stops early. The second test checks that case.

Counting takes O(n). Building the heap from u distinct numbers takes O(u), and k pops take O(k log u).

The first test sorts the result before comparing it. The order of numbers with equal counts depends on the map’s iteration order. That order changes between runs, so a test must not depend on it.

5.5 The k most frequent words, with a heap of size k

The first version puts every distinct number in the heap. If the input has millions of distinct words and you only want the top 3, most of that heap is wasted. The second version never lets the heap grow past k + 1 entries.

The idea. Use a min-heap. Push each (count, word) pair. Whenever the heap has more than k entries, pop once. A min-heap pops its smallest entry, which is the least frequent word seen so far, so the k most frequent words survive (figure 5.4).

The program first turns the text into words, then counts them, then runs the bounded heap.

Listing 5.2 The function (lines 1 to 41). src/bin/top_k_frequent_words.rs

use std::{
    cmp::Reverse as Rev,
    collections::{BinaryHeap, HashMap},
};

/// Returns vec!<(word, frequency)>, most frequent first.
fn top_k_frequent(sentence: &str, k: usize) -> Vec<(String, u32)> {
    // Split string by white spaces. Map it and trim any non alphanumeric characters. Finally,
    // filter all the empty strings and map everything to lowercase characters.
    let words: Vec<String> = sentence
        .split_whitespace()
        .map(|w| w.trim_matches(|c: char| !c.is_alphanumeric()))
        .filter(|w| !w.is_empty())
        .map(str::to_lowercase)
        .collect();

    // Build the frequency map...
    let mut frequencies = HashMap::<&str, u32>::new();
    for word in &words {
        *frequencies
            .entry(word.as_str())
            .or_default() += 1;
    }

    // Keep a min heap of size k. Reverse flips the max heap, so the least
    // frequent word sits on top and gets popped once we go over k.
    let mut min_heap = BinaryHeap::<Rev<(u32, &str)>>::new();
    for (&word, &frequency) in &frequencies {
        min_heap.push(Rev((frequency, word)));
        if min_heap.len() > k {
            min_heap.pop();
        }
    }

    // Ascending order of Reverse means most frequent first.
    min_heap
        .into_sorted_vec()
        .into_iter()
        .map(|Rev((frequency, word))| (word.to_string(), frequency))
        .collect()
}

Turning text into words. The chain splits on whitespace. Then it trims characters that are not letters or digits from both ends of each piece, so "sentence." becomes "sentence". It drops empty pieces, which a lone punctuation mark leaves behind, and lowercases each word. trim_matches takes a closure that decides which characters to trim. map(str::to_lowercase) passes a function by name instead of writing a closure.

Counting without copying. The map is HashMap<&str, u32>. Its keys are references into the words vector, so counting allocates nothing for each occurrence. or_default() inserts the type’s default value, 0 for u32, when the key is missing.

The bounded heap. use std::cmp::Reverse as Rev; gives Reverse a shorter name. Each entry is Rev((frequency, word)), so the heap’s top is the smallest (frequency, word) pair. After each push, if the heap has more than k entries, one pop removes that smallest pair.

The result. into_sorted_vec() returns the entries in ascending order of Rev, which is descending order of frequency. The final map unwraps each Rev with a pattern and copies the word into an owned String. The result must not borrow from words, because words is dropped when the function returns.

The heap holds at most k + 1 entries, so each push and pop costs O(log k). For u distinct words, the heap work is O(u log k) instead of O(u log u), and its memory is O(k) instead of O(u).

5.5.1 The complete program

Listing 5.3 The complete program, with its tests. src/bin/top_k_frequent_words.rs

use std::{
    cmp::Reverse as Rev,
    collections::{BinaryHeap, HashMap},
};

/// Returns vec!<(word, frequency)>, most frequent first.
fn top_k_frequent(sentence: &str, k: usize) -> Vec<(String, u32)> {
    // Split string by white spaces. Map it and trim any non alphanumeric characters. Finally,
    // filter all the empty strings and map everything to lowercase characters.
    let words: Vec<String> = sentence
        .split_whitespace()
        .map(|w| w.trim_matches(|c: char| !c.is_alphanumeric()))
        .filter(|w| !w.is_empty())
        .map(str::to_lowercase)
        .collect();

    // Build the frequency map...
    let mut frequencies = HashMap::<&str, u32>::new();
    for word in &words {
        *frequencies
            .entry(word.as_str())
            .or_default() += 1;
    }

    // Keep a min heap of size k. Reverse flips the max heap, so the least
    // frequent word sits on top and gets popped once we go over k.
    let mut min_heap = BinaryHeap::<Rev<(u32, &str)>>::new();
    for (&word, &frequency) in &frequencies {
        min_heap.push(Rev((frequency, word)));
        if min_heap.len() > k {
            min_heap.pop();
        }
    }

    // Ascending order of Reverse means most frequent first.
    min_heap
        .into_sorted_vec()
        .into_iter()
        .map(|Rev((frequency, word))| (word.to_string(), frequency))
        .collect()
}

fn main() {
    let sentence = r"This is a short sentence. We'll count words.
    Let us count words in this sentence
    ";

    let top_k = top_k_frequent(sentence, 3);
    println!("Top 3 frequency [(word, frequency)] : \n {:#?}", top_k);
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn picks_most_frequent() {
        let result = top_k_frequent("a a a b b c", 2);
        assert_eq!(result, vec![("a".to_string(), 3), ("b".to_string(), 2)]);
    }

    #[test]
    fn ignores_case_and_punctuation() {
        let result = top_k_frequent("Hello, hello! HELLO world.", 1);
        assert_eq!(result, vec![("hello".to_string(), 3)]);
    }

    #[test]
    fn k_larger_than_word_count() {
        assert_eq!(top_k_frequent("x y", 10).len(), 2);
    }

    #[test]
    fn empty_input() {
        assert!(top_k_frequent("", 3).is_empty());
    }
}

The sentence in main is a raw string literal, written r"...". A raw string treats backslashes as ordinary characters, and this one also spans several lines. The tests cover the basic case, uppercase and punctuation, a k larger than the number of words, and an empty input.

$ cargo run --bin top_k_frequent_words
Top 3 frequency [(word, frequency)] :
 [
    (
        "words",
        2,
    ),
    (
        "this",
        2,
    ),
    (
        "sentence",
        2,
    ),
]

Four words occur twice in the sentence: “this”, “sentence”, “count”, and “words”. Only three fit, and “count” was evicted. When two entries have the same count, the tuple comparison falls to the word, and the min-heap evicts the alphabetically smallest one. So among equal counts, this version keeps the words that come last in alphabetical order.

5.6 The median of a stream

The problem. Numbers arrive one at a time. After each one, report the median of all the numbers so far.

The simple idea. Keep every number in a sorted Vec, and read the middle. Inserting into the middle of a sorted vector shifts the elements after it, which costs O(n) per number.

The idea with heaps. Split the numbers into two halves (figure 5.5):

• low, a max-heap holding the smaller half. Its top is the largest of the small numbers.
• high, a min-heap holding the larger half. Its top is the smallest of the large numbers.

Keep two rules. Every number in low is at most every number in high. And low has either the same number of items as high, or one more. Then the median is always at the tops. With an odd count, it is the top of low. With an even count, it is the average of the two tops.

5.6.1 Adding a number

Listing 5.4 The structure and add_num (lines 1 to 27). src/bin/median_finder.rs

use std::{cmp::Reverse, collections::BinaryHeap};

#[derive(Default)]
pub struct MedianFinder {
    low: BinaryHeap<i32>,
    high: BinaryHeap<Reverse<i32>>,
}

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

    pub fn add_num(&mut self, num: i32) {
        self.low.push(num);

        if let Some(val) = self.low.pop() {
            self.high.push(Reverse(val));
        }

        if self.high.len() > self.low.len()
            && let Some(Reverse(val)) = self.high.pop()
        {
            self.low.push(val);
        }
    }
    // ...
}

add_num keeps both rules with three moves, and never compares numbers itself:

1. Push the new number into low.
2. Pop the largest number from low and push it into high. This keeps the first rule: whatever high receives is at least as large as everything left in low.
3. If high now has more items than low, pop its smallest and push it back into low. This keeps the second rule.

Trace the first three numbers of the sample, 6, 10, and 2:

Each move is one heap operation, so adding a number is O(log n).

5.6.2 Reading the median

Listing 5.5 find_median (lines 28 to 37).

impl MedianFinder {
    // ...
    pub fn find_median(&self) -> Option<f64> {
        let max_low = *self.low.peek()?;

        if self.low.len() > self.high.len() {
            Some(max_low as f64)
        } else {
            let Reverse(min_high) = *self.high.peek()?;
            Some((max_low as f64 + min_high as f64) / 2.0)
        }
    }
}

*self.low.peek()? reads the top of low. The ? works on Option the same way it works on Result. If no number has been added, peek returns None, and so does the function. The function returns Option<f64>, a floating-point number, because the average of two integers can be a half.

let Reverse(min_high) = *self.high.peek()?; unwraps the Reverse in the let pattern itself.

5.6.3 The complete program

Listing 5.6 The complete program, with its tests. src/bin/median_finder.rs

use std::{cmp::Reverse, collections::BinaryHeap};

#[derive(Default)]
pub struct MedianFinder {
    low: BinaryHeap<i32>,
    high: BinaryHeap<Reverse<i32>>,
}

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

    pub fn add_num(&mut self, num: i32) {
        self.low.push(num);

        if let Some(val) = self.low.pop() {
            self.high.push(Reverse(val));
        }

        if self.high.len() > self.low.len()
            && let Some(Reverse(val)) = self.high.pop()
        {
            self.low.push(val);
        }
    }

    pub fn find_median(&self) -> Option<f64> {
        let max_low = *self.low.peek()?;

        if self.low.len() > self.high.len() {
            Some(max_low as f64)
        } else {
            let Reverse(min_high) = *self.high.peek()?;
            Some((max_low as f64 + min_high as f64) / 2.0)
        }
    }
}

const SAMPLE: [i32; 11] = [6, 10, 2, 6, 5, 0, 6, 3, 1, 0, 0];

fn main() {
    let mut finder = MedianFinder::new();

    for value in SAMPLE {
        finder.add_num(value);
        match finder.find_median() {
            Some(median) => println!("after {value:>2}: median {median}"),
            None => println!("after {value:>2}: no median"),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// The median of a sorted slice, computed directly.
    fn median_of_sorted(values: &[i32]) -> f64 {
        let middle = values.len() / 2;
        if values.len() % 2 == 1 {
            values[middle] as f64
        } else {
            (values[middle - 1] as f64 + values[middle] as f64) / 2.0
        }
    }

    #[test]
    fn an_empty_finder_has_no_median() {
        assert_eq!(MedianFinder::new().find_median(), None);
    }

    #[test]
    fn one_value_is_the_median() {
        let mut finder = MedianFinder::new();
        finder.add_num(7);
        assert_eq!(finder.find_median(), Some(7.0));
    }

    #[test]
    fn an_even_count_averages_the_middle_two() {
        let mut finder = MedianFinder::new();
        finder.add_num(1);
        finder.add_num(2);
        assert_eq!(finder.find_median(), Some(1.5));
    }

    #[test]
    fn three_values_are_ordered() {
        let mut finder = MedianFinder::new();
        for (value, median) in [1, 2, 3]
            .into_iter()
            .zip([1.0, 1.5, 2.0])
        {
            finder.add_num(value);
            assert_eq!(finder.find_median(), Some(median));
        }
    }

    #[test]
    fn duplicates_are_kept() {
        let mut finder = MedianFinder::new();
        for value in [4, 4, 4, 4] {
            finder.add_num(value);
        }
        assert_eq!(finder.find_median(), Some(4.0));
    }

    #[test]
    fn the_stream_matches_a_sort_on_every_step() {
        let mut finder = MedianFinder::new();
        let mut seen = Vec::new();

        for value in SAMPLE {
            finder.add_num(value);
            seen.push(value);
            seen.sort_unstable();

            assert_eq!(finder.find_median(), Some(median_of_sorted(&seen)));
        }
    }

    #[test]
    fn the_two_heaps_never_differ_by_more_than_one() {
        let mut finder = MedianFinder::new();

        for value in SAMPLE {
            finder.add_num(value);
            assert!(
                finder
                    .low
                    .len()
                    .abs_diff(finder.high.len())
                    <= 1
            );
        }
    }
}

main adds the eleven sample values and prints the median after each one.

$ cargo run --bin median_finder
after  6: median 6
after 10: median 8
after  2: median 6
after  6: median 6
after  5: median 6
after  0: median 5.5
after  6: median 6
after  3: median 5.5
after  1: median 5
after  0: median 4
after  0: median 3

Look at the test the_stream_matches_a_sort_on_every_step. After every added number, it sorts a copy of all the numbers so far. Then it computes the median the slow way, with the helper median_of_sorted. Then it checks that the heaps give the same answer. A test like this, which compares a fast method with a slow method you can trust, is called an oracle test. It catches mistakes that a few hand-picked examples miss.

The last test checks the second rule directly: after every number, the two heaps’ sizes differ by at most one.

Chapter 6 covers recursion and dynamic programming. You will solve problems by combining answers to smaller versions of themselves, and save those answers so none is computed twice.

Exercises

1. Implement a max-heap yourself on a Vec<i32>, with push, pop, and peek, using the index formulas from section 5.1. Test it against BinaryHeap on random input.
2. Change top_k_frequent_words to keep the alphabetically first words among ties, as the tip in section 5.5 describes, and add a test.
3. Add a method len to MedianFinder that returns how many numbers it holds.
4. Extend the running median to remove numbers as well, so it can report the median of the last 100 numbers only. Hint: remove lazily, by remembering which numbers are waiting to be deleted.