362 Intermediate

362: Trie Structure

Functional Programming

Tutorial Video

Text description (accessibility)

This video demonstrates the "362: Trie Structure" functional Rust example. Difficulty level: Intermediate. Key concepts covered: Functional Programming. Hash maps give O(1) exact key lookup but can't answer prefix queries: "list all words starting with 'pre'" requires scanning all keys. Key difference from OCaml: | Aspect | Rust `HashMap<char, TrieNode>` | OCaml `CharMap.t trie` |

Tutorial

The Problem

Hash maps give O(1) exact key lookup but can't answer prefix queries: "list all words starting with 'pre'" requires scanning all keys. A trie (retrieval tree, Fredkin 1960) stores strings by decomposing them into characters — each node represents one character, paths from root to end-marked nodes spell out stored words. Lookup is O(m) where m is the key length, independent of how many words are stored. Tries power autocomplete in IDEs and search engines, IP routing tables (compact trie on bit prefixes), spell checkers, and dictionary compression (compressed tries / DAWG). They're the data structure behind DNS resolution caches.

🎯 Learning Outcomes

• Implement a TrieNode with children: HashMap<char, TrieNode> and is_end: bool

• Insert a word by traversing/creating one node per character

• Search for exact words by traversing the path and checking is_end

• Find all words with a given prefix using recursive subtree traversal (DFS)

• Understand why lookup is O(m) regardless of vocabulary size

• Compare trie prefix queries to linear scan over a sorted list

Code Example

use std::collections::HashMap;

struct TrieNode {
    children: HashMap<char, TrieNode>,
    is_end: bool,
}

module CharMap = Map.Make(Char)

type trie = { is_end: bool; children: trie CharMap.t }
let empty = { is_end=false; children=CharMap.empty }

Key Differences

Aspect	Rust `HashMap<char, TrieNode>`	OCaml `CharMap.t trie`
Children map	`HashMap` (O(1) lookup)	`Map.Make(Char)` (O(log σ))
Mutability	In-place (`&mut self`)	Persistent (new trie per insert)
Alphabet size	Any (unicode chars)	Any (via functor)
Memory sharing	None — copy on clone	Structural sharing
Compact trie	Manual or `radix-trie` crate	Manual implementation

OCaml Approach

OCaml's recursive types with Map children:

module CharMap = Map.Make(Char)

type trie = {
  children: trie CharMap.t;
  is_end: bool;
}

let empty = { children = CharMap.empty; is_end = false }

let rec insert node = function
  | [] -> { node with is_end = true }
  | c :: rest ->
    let child = try CharMap.find c node.children with Not_found -> empty in
    let updated = insert child rest in
    { node with children = CharMap.add c updated node.children }

let insert_word trie word =
  insert trie (List.of_seq (String.to_seq word))

OCaml's functional trie returns new nodes on insert (persistent). CharMap replaces HashMap<char, TrieNode> with a functional balanced BST. The structure is immutable — each insert returns a new trie sharing unchanged subtrees with the original.

Full Source

#![allow(clippy::all)]
//! Trie (Prefix Tree) for efficient string lookups
//!
//! O(m) lookup and prefix search where m is the key length.

use std::collections::HashMap;

/// A trie node containing children and an end-of-word marker
#[derive(Default, Debug, Clone)]
pub struct TrieNode {
    children: HashMap<char, TrieNode>,
    is_end: bool,
}

impl TrieNode {
    /// Create a new empty trie node
    pub fn new() -> Self {
        Self::default()
    }
}

/// A trie (prefix tree) for string storage and lookup
#[derive(Debug, Clone)]
pub struct Trie {
    root: TrieNode,
}

impl Trie {
    // === Approach 1: Basic insert/search API ===

    /// Create a new empty trie
    pub fn new() -> Self {
        Self {
            root: TrieNode::default(),
        }
    }

    /// Insert a word into the trie - O(m)
    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_end = true;
    }

    /// Search for an exact word - O(m)
    pub fn search(&self, word: &str) -> bool {
        let mut node = &self.root;
        for c in word.chars() {
            match node.children.get(&c) {
                None => return false,
                Some(n) => node = n,
            }
        }
        node.is_end
    }

    /// Check if any word starts with the given prefix - O(m)
    pub fn starts_with(&self, prefix: &str) -> bool {
        let mut node = &self.root;
        for c in prefix.chars() {
            match node.children.get(&c) {
                None => return false,
                Some(n) => node = n,
            }
        }
        true
    }

    // === Approach 2: Prefix collection ===

    /// Get all words with a given prefix - O(m + k) where k is result count
    pub fn words_with_prefix(&self, prefix: &str) -> Vec<String> {
        let mut node = &self.root;
        for c in prefix.chars() {
            match node.children.get(&c) {
                None => return vec![],
                Some(n) => node = n,
            }
        }
        let mut result = Vec::new();
        Self::collect(node, prefix.to_string(), &mut result);
        result
    }

    fn collect(node: &TrieNode, current: String, result: &mut Vec<String>) {
        if node.is_end {
            result.push(current.clone());
        }
        for (c, child) in &node.children {
            let mut next = current.clone();
            next.push(*c);
            Self::collect(child, next, result);
        }
    }

    // === Approach 3: Additional utilities ===

    /// Remove a word from the trie (marks it as not an end)
    pub fn remove(&mut self, word: &str) -> bool {
        let mut node = &mut self.root;
        for c in word.chars() {
            match node.children.get_mut(&c) {
                None => return false,
                Some(n) => node = n,
            }
        }
        if node.is_end {
            node.is_end = false;
            true
        } else {
            false
        }
    }

    /// Count all words in the trie
    pub fn count_words(&self) -> usize {
        Self::count_words_recursive(&self.root)
    }

    fn count_words_recursive(node: &TrieNode) -> usize {
        let mut count = if node.is_end { 1 } else { 0 };
        for child in node.children.values() {
            count += Self::count_words_recursive(child);
        }
        count
    }

    /// Get all words in the trie
    pub fn all_words(&self) -> Vec<String> {
        let mut result = Vec::new();
        Self::collect(&self.root, String::new(), &mut result);
        result
    }

    /// Check if the trie is empty
    pub fn is_empty(&self) -> bool {
        self.root.children.is_empty() && !self.root.is_end
    }

    /// Count words with a given prefix
    pub fn count_with_prefix(&self, prefix: &str) -> usize {
        let mut node = &self.root;
        for c in prefix.chars() {
            match node.children.get(&c) {
                None => return 0,
                Some(n) => node = n,
            }
        }
        Self::count_words_recursive(node)
    }
}

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

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

    #[test]
    fn test_insert_and_search() {
        let mut trie = Trie::new();
        trie.insert("hello");
        trie.insert("world");

        assert!(trie.search("hello"));
        assert!(trie.search("world"));
        assert!(!trie.search("hell"));
        assert!(!trie.search("worlds"));
    }

    #[test]
    fn test_starts_with() {
        let mut trie = Trie::new();
        trie.insert("apple");
        trie.insert("application");

        assert!(trie.starts_with("app"));
        assert!(trie.starts_with("apple"));
        assert!(trie.starts_with("appli"));
        assert!(!trie.starts_with("xyz"));
    }

    #[test]
    fn test_words_with_prefix() {
        let mut trie = Trie::new();
        for word in ["cat", "car", "card", "care", "cart"] {
            trie.insert(word);
        }

        let mut results = trie.words_with_prefix("car");
        results.sort();
        assert_eq!(results, vec!["car", "card", "care", "cart"]);
    }

    #[test]
    fn test_empty_prefix() {
        let mut trie = Trie::new();
        trie.insert("a");
        trie.insert("b");

        let mut all = trie.words_with_prefix("");
        all.sort();
        assert_eq!(all, vec!["a", "b"]);
    }

    #[test]
    fn test_remove() {
        let mut trie = Trie::new();
        trie.insert("hello");
        trie.insert("help");

        assert!(trie.search("hello"));
        assert!(trie.remove("hello"));
        assert!(!trie.search("hello"));
        assert!(trie.search("help"));
    }

    #[test]
    fn test_count_words() {
        let mut trie = Trie::new();
        assert_eq!(trie.count_words(), 0);

        trie.insert("a");
        trie.insert("ab");
        trie.insert("abc");
        assert_eq!(trie.count_words(), 3);
    }

    #[test]
    fn test_count_with_prefix() {
        let mut trie = Trie::new();
        for word in ["apple", "app", "application", "apply", "banana"] {
            trie.insert(word);
        }

        assert_eq!(trie.count_with_prefix("app"), 4);
        assert_eq!(trie.count_with_prefix("ban"), 1);
        assert_eq!(trie.count_with_prefix("xyz"), 0);
    }

    #[test]
    fn test_all_words() {
        let mut trie = Trie::new();
        trie.insert("rust");
        trie.insert("ruby");

        let mut words = trie.all_words();
        words.sort();
        assert_eq!(words, vec!["ruby", "rust"]);
    }

    #[test]
    fn test_is_empty() {
        let mut trie = Trie::new();
        assert!(trie.is_empty());

        trie.insert("test");
        assert!(!trie.is_empty());
    }

    #[test]
    fn test_unicode_support() {
        let mut trie = Trie::new();
        trie.insert("日本語");
        trie.insert("日本");

        assert!(trie.search("日本語"));
        assert!(trie.search("日本"));
        assert!(trie.starts_with("日"));
    }
}

(* OCaml: simple trie *)

module CharMap = Map.Make(Char)

type trie = { is_end: bool; children: trie CharMap.t }
let empty = { is_end=false; children=CharMap.empty }

let insert t word =
  let n = String.length word in
  let rec go node i =
    if i = n then { node with is_end=true }
    else
      let c = word.[i] in
      let child = try CharMap.find c node.children with Not_found -> empty in
      let new_child = go child (i+1) in
      { node with children=CharMap.add c new_child node.children }
  in go t 0

let search t word =
  let n = String.length word in
  let rec go node i =
    if i=n then node.is_end
    else match CharMap.find_opt word.[i] node.children with
    | None -> false | Some child -> go child (i+1)
  in go t 0

let () =
  let t = List.fold_left insert empty ["apple";"app";"application";"apply";"banana"] in
  List.iter (fun w -> Printf.printf "%s: %b\n" w (search t w))
    ["apple";"app";"ap";"application";"ban";"banana";"cat"]

✓ Tests Rust test suite

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

    #[test]
    fn test_insert_and_search() {
        let mut trie = Trie::new();
        trie.insert("hello");
        trie.insert("world");

        assert!(trie.search("hello"));
        assert!(trie.search("world"));
        assert!(!trie.search("hell"));
        assert!(!trie.search("worlds"));
    }

    #[test]
    fn test_starts_with() {
        let mut trie = Trie::new();
        trie.insert("apple");
        trie.insert("application");

        assert!(trie.starts_with("app"));
        assert!(trie.starts_with("apple"));
        assert!(trie.starts_with("appli"));
        assert!(!trie.starts_with("xyz"));
    }

    #[test]
    fn test_words_with_prefix() {
        let mut trie = Trie::new();
        for word in ["cat", "car", "card", "care", "cart"] {
            trie.insert(word);
        }

        let mut results = trie.words_with_prefix("car");
        results.sort();
        assert_eq!(results, vec!["car", "card", "care", "cart"]);
    }

    #[test]
    fn test_empty_prefix() {
        let mut trie = Trie::new();
        trie.insert("a");
        trie.insert("b");

        let mut all = trie.words_with_prefix("");
        all.sort();
        assert_eq!(all, vec!["a", "b"]);
    }

    #[test]
    fn test_remove() {
        let mut trie = Trie::new();
        trie.insert("hello");
        trie.insert("help");

        assert!(trie.search("hello"));
        assert!(trie.remove("hello"));
        assert!(!trie.search("hello"));
        assert!(trie.search("help"));
    }

    #[test]
    fn test_count_words() {
        let mut trie = Trie::new();
        assert_eq!(trie.count_words(), 0);

        trie.insert("a");
        trie.insert("ab");
        trie.insert("abc");
        assert_eq!(trie.count_words(), 3);
    }

    #[test]
    fn test_count_with_prefix() {
        let mut trie = Trie::new();
        for word in ["apple", "app", "application", "apply", "banana"] {
            trie.insert(word);
        }

        assert_eq!(trie.count_with_prefix("app"), 4);
        assert_eq!(trie.count_with_prefix("ban"), 1);
        assert_eq!(trie.count_with_prefix("xyz"), 0);
    }

    #[test]
    fn test_all_words() {
        let mut trie = Trie::new();
        trie.insert("rust");
        trie.insert("ruby");

        let mut words = trie.all_words();
        words.sort();
        assert_eq!(words, vec!["ruby", "rust"]);
    }

    #[test]
    fn test_is_empty() {
        let mut trie = Trie::new();
        assert!(trie.is_empty());

        trie.insert("test");
        assert!(!trie.is_empty());
    }

    #[test]
    fn test_unicode_support() {
        let mut trie = Trie::new();
        trie.insert("日本語");
        trie.insert("日本");

        assert!(trie.search("日本語"));
        assert!(trie.search("日本"));
        assert!(trie.starts_with("日"));
    }
}

Deep Comparison

OCaml vs Rust: Trie (Prefix Tree)

Side-by-Side Comparison

Type Definition

OCaml:

module CharMap = Map.Make(Char)

type trie = { is_end: bool; children: trie CharMap.t }
let empty = { is_end=false; children=CharMap.empty }

Rust:

use std::collections::HashMap;

struct TrieNode {
    children: HashMap<char, TrieNode>,
    is_end: bool,
}

Insert Operation

OCaml:

let insert t word =
  let n = String.length word in
  let rec go node i =
    if i = n then { node with is_end=true }
    else
      let c = word.[i] in
      let child = try CharMap.find c node.children with Not_found -> empty in
      let new_child = go child (i+1) in
      { node with children=CharMap.add c new_child node.children }
  in go t 0

Rust:

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_end = true;
}

Search Operation

OCaml:

let search t word =
  let n = String.length word in
  let rec go node i =
    if i=n then node.is_end
    else match CharMap.find_opt word.[i] node.children with
    | None -> false
    | Some child -> go child (i+1)
  in go t 0

Rust:

fn search(&self, word: &str) -> bool {
    let mut node = &self.root;
    for c in word.chars() {
        match node.children.get(&c) {
            None => return false,
            Some(n) => node = n,
        }
    }
    node.is_end
}

Key Differences

Aspect	OCaml	Rust
Map type	`CharMap` (functor)	`HashMap<char, _>`
Iteration style	Recursive with index	Iterative with `for`
Not found handling	Exception or `find_opt`	`Option` via `get`
Mutability	Functional (new nodes)	In-place mutation
Entry API	Manual insert	`entry().or_default()`

Memory Model

OCaml: Creates new nodes on each insert (functional style). The old trie is preserved (persistent data structure).

Rust: Mutates nodes in place. Previous state is lost unless explicitly cloned.

Performance Characteristics

Operation	OCaml	Rust
Insert	O(m) - creates new path	O(m) - mutates in place
Search	O(m)	O(m)
Prefix search	O(m + k)	O(m + k)
Memory	More (path copying)	Less (in-place)

Where m = word length, k = number of results

Exercises

Delete a word: Implement remove(&mut self, word: &str) that marks is_end = false for the word's terminal node; also prune empty subtrees (nodes with no children and is_end = false).

Longest common prefix: Given a list of words inserted into a trie, find the longest common prefix by traversing the root while each node has exactly one child and is_end = false.

Compressed trie: Implement a Patricia trie where edge labels are substrings (not single chars), reducing the number of nodes from O(total_chars) to O(words) for strings with long common prefixes.

Open Source Repos

functional-rust

View the source for this example on GitHub — OCaml and Rust side by side in the repo.

Rust