Tries in Rust

Open the messages app on your phone and type "ca". Before you finish the word, it offers "cat", "car" and "cap". It did not read through every word in the dictionary to find them. It went straight to them.

The trick is a structure called a trie (said like "try", from the middle of the word re-trie-val). It is a tree where each step down is one letter. A word is a path from the top.

A prefix is the start of a word. "ca" is a prefix of "cat". A trie is built around prefixes: words that begin the same way share the same path for as long as they agree, and split only where they differ.

Reach for a trie when you ask questions about the start of words: autocomplete, spell checking, "does any word begin with this?", word-game helpers, and search boxes that suggest as you type.

The picture holds five words: cap, car, cat, do and dog. The three "ca" words share two steps, down through c and then a. Only the last letter differs, so the path forks there.

Now search for the prefix "ca". Start at the top, step to c, step to a, and you are there: two steps, one per letter. Everything below that node is a word that starts with "ca". The suggestions are simply what hangs underneath.

Look at how few steps that took. It does not matter whether the trie holds five words or five million. The number of steps depends only on the length of the prefix. That is the superpower of a trie.

There is one wrinkle. "do" is a word, and it is also the start of "dog". So reaching a node does not tell you a word ends there. Each node needs a flag, which the picture shows as a ring around the node. A ringed node means "a word ends here".

Inside, every node holds two things: a map from letter to the next node, and that flag. A map holding smaller maps holding smaller maps is why people call this a structure of nested maps.

Remember: a trie's speed depends on how long the word is, not on how many words it holds.

use std::collections::BTreeMap;

/// One step in a word. It says which letters can come next, and whether a
/// word ends right here.
#[derive(Default)]
pub struct TrieNode {
    pub children: BTreeMap<char, TrieNode>,
    pub is_word: bool,
}

/// A trie: a tree of letters, where each word is a path from the root.
#[derive(Default)]
pub struct Trie {
    root: TrieNode,
}

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

    pub fn root(&self) -> &TrieNode {
        &self.root
    }

    /// Walks down the word, adding a node for each letter that is missing.
    pub fn insert(&mut self, word: &str) {
        let mut node = &mut self.root;
        for c in word.chars() {
            node = node.children.entry(c).or_default();
        }
        node.is_word = true;
    }

    /// Follows `prefix` down from the root. `None` means the path breaks off.
    fn find(&self, prefix: &str) -> Option<&TrieNode> {
        let mut node = &self.root;
        for c in prefix.chars() {
            node = node.children.get(&c)?;
        }
        Some(node)
    }

    /// True only for a whole word that was inserted.
    pub fn contains(&self, word: &str) -> bool {
        self.find(word).is_some_and(|node| node.is_word)
    }

    /// True if any inserted word begins with `prefix`.
    pub fn starts_with(&self, prefix: &str) -> bool {
        self.find(prefix).is_some()
    }
}

Go through it in order.

• TrieNode is one step. children maps a letter to the next TrieNode, and is_word is the ring from the picture. #[derive(Default)] lets Rust build an empty node on its own, with no children and false.
• It uses a BTreeMap from lesson 9, so the letters come back in alphabetical order. That is why the suggestions later arrive already sorted.
• A TrieNode holds a map of TrieNodes, which sounds like a type that contains itself forever. It works without a Box because a map keeps its contents on the heap, so the node itself has a fixed size. Lesson 7 needed Box because it held a child directly.
• insert walks down the word one letter at a time. entry(c).or_default() means "go to the child for c, and if there is none, make an empty one". After the last letter, it sets is_word to true.
• find follows a prefix down from the root. The ? returns None the moment a letter has no child, meaning the path breaks off.
• contains and starts_with both use find, and differ in one detail. contains also checks the ring. starts_with only needs the path to exist.

let mut trie = Trie::new();
for word in ["cat", "car", "cap", "dog", "do"] {
    trie.insert(word);
}

assert!(trie.contains("car"));
assert!(!trie.contains("ca")); // a prefix, but not a word
assert!(trie.starts_with("ca"));
assert!(!trie.starts_with("cb"));
assert!(trie.contains("do") && trie.contains("dog"));
println!("{}", trie.contains("cat"));

"ca" is not a word, but it is a prefix, and the trie tells the two apart. The Playground link on the code above runs this.

Now the autocomplete itself. To list every word under a prefix, walk to the prefix node and then collect everything below it.

impl Trie {
    /// Every stored word that begins with `prefix`, in alphabetical order.
    pub fn suggestions(&self, prefix: &str) -> Vec<String> {
        let mut found = Vec::new();
        if let Some(node) = self.find(prefix) {
            let mut word = prefix.to_string();
            collect_words(node, &mut word, &mut found);
        }
        found
    }
}

fn collect_words(node: &TrieNode, word: &mut String, found: &mut Vec<String>) {
    if node.is_word {
        found.push(word.clone());
    }
    // A BTreeMap hands its children back in letter order, so the words do too.
    for (&c, child) in &node.children {
        word.push(c);
        collect_words(child, word, found);
        word.pop(); // undo the letter, so the next sibling starts clean
    }
}

collect_words is recursive. At each node it first checks the ring: if a word ends here, it saves a copy of the letters so far. Then, for each child, it adds that child's letter to word, calls itself, and takes the letter off again with word.pop(). That last step is like retracing your steps: it puts word back as it was, so the next sibling starts from the right place. Because the children are in letter order, the words are found in alphabetical order.

Remember: add a letter on the way down, and remove it on the way back up.

The standard library has no trie. But if all you need is "give me the words that start with this", a sorted set can do it, using what you learned in lesson 9.

use std::collections::BTreeSet;

let words: BTreeSet<&str> = ["cap", "car", "cat", "do", "dog"].into_iter().collect();
let prefix = "ca";

// Words that share a prefix sit side by side in sorted order. Jump to
// where the prefix would go, then stop at the first word that differs.
let found: Vec<&str> = words
    .range(prefix..)
    .take_while(|word| word.starts_with(prefix))
    .copied()
    .collect();
println!("{:?}", found);
assert_eq!(found, ["cap", "car", "cat"]);

In sorted order, all the words that start with "ca" sit side by side. So range(prefix..) jumps to the place where "ca" would go, and take_while keeps going until it meets a word that no longer starts with it. That is a range query with a stopping rule.

So which should you use? The sorted set is shorter and often good enough. A hand-built trie earns its keep when you need to step through the letters one at a time, attach information to each prefix, or match patterns, as in problem 3. Each lookup in the set compares whole strings, and in a trie it only follows letters.

One more option, when words are only lowercase a to z: give each node an array of 26 slots instead of a map. It is faster, and it wastes memory when most slots are empty.

None of those depend on how many words are stored. That is the strength. The BTreeMap inside each node adds a tiny extra cost to pick the right letter, but with at most a few dozen letters to choose from, it is safe to ignore. The price you pay is memory: a trie of words that share little can use far more space than the plain words would.

These use the Trie you just built. The Playground link on each starter includes it, along with a small demo. Each starter has a todo!(), a placeholder that panics if it runs. Replace it with your code.

pub fn count_words(node: &TrieNode) -> usize {
    todo!("how many words end at or below this node?")
}

pub fn count_words(node: &TrieNode) -> usize {
    // This node counts once if a word ends here. Then add each child's total.
    let here = usize::from(node.is_word);
    let below: usize = node.children.values().map(count_words).sum();
    here + below
}

pub fn count_with_prefix(trie: &Trie, prefix: &str) -> usize {
    todo!("how many stored words begin with the prefix?")
}

pub fn count_with_prefix(trie: &Trie, prefix: &str) -> usize {
    let mut node = trie.root();
    for c in prefix.chars() {
        match node.children.get(&c) {
            Some(next) => node = next,
            None => return 0, // the path breaks off, so nothing starts this way
        }
    }
    // Everything below this node starts with the prefix.
    count_words(node)
}

pub fn matches(trie: &Trie, pattern: &str) -> bool {
    todo!("does any stored word fit the pattern, where . is any letter?")
}

pub fn matches(trie: &Trie, pattern: &str) -> bool {
    fn walk(node: &TrieNode, pattern: &[char]) -> bool {
        let Some((&first, rest)) = pattern.split_first() else {
            return node.is_word; // the pattern is used up: does a word end here?
        };
        if first == '.' {
            // A dot could be any letter, so try every branch.
            node.children.values().any(|child| walk(child, rest))
        } else {
            node.children
                .get(&first)
                .is_some_and(|child| walk(child, rest))
        }
    }

    let chars: Vec<char> = pattern.chars().collect();
    walk(trie.root(), &chars)
}

Trie questions almost always involve prefixes. A few habits help.

• Spot the signals. "Starts with", "autocomplete", "dictionary of words", "longest common prefix", and "word search" are all trie territory.
• Draw it first. Sketch three or four words as a tree on the board. It makes the design obvious, and it shows the interviewer you can explain the idea, not just code it.
• Remember the ring. The most common bug is forgetting the end-of-word flag. "do" and "dog" both being words is the example to test with.
• Pick the node's map on purpose. A HashMap<char, Node> is fast and unordered. A BTreeMap gives sorted output for free. An array of 26 is quickest for lowercase letters only. Naming the trade-off is worth more than the choice.
• Know the alternative. For a simple prefix question, a sorted list or a BTreeSet with a range query can be enough. Saying so shows you do not reach for a trie out of habit.
• Say the cost. "Insert and search are O(m) in the length of the word, whatever the number of words" is a complete answer. Add that the trade-off is memory.

Made by VeraStack Labs