Step 3 of 6

Memoization with a HashMap

The slow Fibonacci recomputes the same values again and again. Memoization fixes that: the first time you compute an answer for some input, store it in a HashMap (the memo); every later call with that input just looks it up. Now each different input is computed only once.

use std::collections::HashMap;

fn fib(n: u64, memo: &mut HashMap<u64, u64>) -> u64 {
    if n < 2 {
        return n;
    }
    if let Some(&known) = memo.get(&n) {
        return known;
    }
    let value = fib(n - 1, memo) + fib(n - 2, memo);
    memo.insert(n, value);
    value
}

fn main() {
    let mut memo = HashMap::new();
    println!("fib(50) = {}", fib(50, &mut memo));
    println!("fib(90) = {}", fib(90, &mut memo));
    println!("memo holds {} answers", memo.len());
}

fib(50) = 12586269025
fib(90) = 2880067194370816120
memo holds 89 answers

fib(90) the slow way would take about 10¹⁹ calls. With the memo it's 89 entries and a couple of hundred calls: O(n).

This is dynamic programming (DP), top-down style. It applies when a problem has:

• overlapping subproblems: the same smaller questions come up many times, and
• optimal substructure: the answer can be built from the answers to smaller questions.

Passing the memo around

Rust has no convenient global mutable variables, and that's fine: pass the memo down as &mut HashMap. A common shape is a public function that creates the memo and a helper that does the work:

fn count(n: u64) -> u64 {
    count_with(n, &mut HashMap::new())
}

The key must describe the whole subproblem. If the answer depends on a position in a grid, the key is the position: HashMap<(usize, usize), u64>. Tuples of Hash + Eq types work as keys.

Common mistakes

• Checking the memo after recursing. Look up first, compute only on a miss, store before returning.
• A new HashMap inside the recursive function. Each call gets an empty memo and nothing is ever reused.
• entry().or_insert_with with a recursive closure. It looks neat but can't compile: the entry holds a mutable borrow of the map while the closure needs another one.

use std::collections::HashMap;

fn fib(n: u64, memo: &mut HashMap<u64, u64>) -> u64 {
    if n < 2 {
        return n;
    }
    *memo.entry(n).or_insert_with(|| fib(n - 1, memo) + fib(n - 2, memo))
}

fn main() {
    println!("{}", fib(10, &mut HashMap::new()));
}

error[E0500]: closure requires unique access to `*memo` but it is already borrowed

Use get, compute, then insert, as in the example.

Your task

A crab climbs a staircase of n steps, taking 1, 2 or 3 steps at a time. In how many different ways can it reach the top? count_ways(4) is 7 (1+1+1+1, 1+1+2, 1+2+1, 2+1+1, 2+2, 1+3, 3+1). There's 1 way to climb 0 steps (do nothing). The number of ways for n is the sum of the ways for n - 1, n - 2 and n - 3 (whichever exist). Write the memoized helper ways_from. The tests ask for n = 70, which takes about 10¹⁸ calls without the memo.

Challenge 2

A delivery robot moves through a warehouse grid from the top-left cell to the bottom-right cell, only ever moving right or down. # cells are shelves it can't enter. Write fn count_paths(grid: &[&str]) -> u64: the number of different routes (0 if the start or the end is a shelf, or the grid is empty).

Recursively, the number of routes from cell (r, c) is the routes from the cell below plus the routes from the cell to the right; from the bottom-right cell it's 1. Memoize on the position with a HashMap<(usize, usize), u64>. The tests include an empty 18 by 18 grid, which has over two billion routes.

Challenge 3

A search engine splits hashtags into words: can "applepenapple" be written as a sequence of dictionary words? Yes: apple + pen + apple. Write fn can_segment(text: &str, words: &[&str]) -> bool (words may be reused; an empty text can always be segmented).

Recursively: the text from position start can be segmented if some dictionary word matches at start and the rest (from start + word.len()) can be segmented. Without a memo, a text like "aaaa...ab" with the words a, aa, aaa, aaaa makes the recursion try an astronomical number of splits. Memoize on start with a HashMap<usize, bool>. Use text[start..].starts_with(word). (The tests use ASCII text, so byte positions are fine.)

Previous: Branching recursion and backtracking Next: Bottom-up DP: filling a table