821: Trie for Autocomplete and Prefix Search

Difficulty: 3 Level: Intermediate Build a prefix tree for O(|prefix|) insert, lookup, and prefix enumeration — the backbone of autocomplete engines.

The Problem This Solves

A trie (prefix tree) stores strings so that all strings sharing a common prefix share the same path from the root. This makes prefix queries — "find all words starting with `com`" — cost O(|prefix| + number of results) rather than O(n × |word|) for a linear scan. Use tries for autocomplete (type-ahead search), IP routing tables (longest prefix match), spell-check suggestion, and dictionary compression. Any system where you need fast prefix lookup over a set of strings benefits from a trie. Unlike a hash map, a trie can enumerate all words with a given prefix efficiently and supports ordered iteration naturally. This implementation supports `insert`, exact `search`, `starts_with` (prefix existence), and `autocomplete` (enumerate all completions for a prefix). It stores complete words at terminal nodes.

The Intuition

Each node represents a character position. The path from the root to a node spells out a prefix. Each node has a map from characters to child nodes, and a flag indicating whether this node ends a complete word. Insert: walk character by character, creating nodes as needed. Mark the final node as a word ending. Search: walk character by character; if any step has no matching child, the word isn't present. Autocomplete: find the node for the prefix, then DFS the subtree collecting all paths to word-ending nodes. O(|word|) for insert/search. O(|prefix| + total_matching_chars) for autocomplete. Space: O(total chars across all words) in the worst case, but shared prefixes save memory proportionally. In OCaml, you'd define a recursive type `type node = { children: (char, node) Hashtbl.t; is_end: bool }`. In Rust, `HashMap<char, Box<TrieNode>>` achieves the same — boxed children to handle the recursive type with known size.

How It Works in Rust

use std::collections::HashMap;

#[derive(Default)]
struct TrieNode {
 children: HashMap<char, Box<TrieNode>>,
 is_end: bool,
}

struct Trie { root: TrieNode }

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

 fn insert(&mut self, word: &str) {
     let mut node = &mut self.root;
     for ch in word.chars() {
         // entry API: insert if absent, then navigate into child
         node = node.children
             .entry(ch)
             .or_insert_with(|| Box::new(TrieNode::default()));
     }
     node.is_end = true;
 }

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

 fn starts_with(&self, prefix: &str) -> bool {
     let mut node = &self.root;
     for ch in prefix.chars() {
         match node.children.get(&ch) {
             Some(child) => node = child,
             None => return false,
         }
     }
     true // reached end of prefix — it exists
 }

 fn autocomplete(&self, prefix: &str) -> Vec<String> {
     // Navigate to prefix node
     let mut node = &self.root;
     for ch in prefix.chars() {
         match node.children.get(&ch) {
             Some(child) => node = child,
             None => return vec![], // prefix not in trie
         }
     }
     // DFS the subtree, collecting completions
     let mut results = vec![];
     Self::collect(node, &mut prefix.to_string(), &mut results);
     results
 }

 fn collect(node: &TrieNode, current: &mut String, results: &mut Vec<String>) {
     if node.is_end { results.push(current.clone()); }
     for (&ch, child) in &node.children {
         current.push(ch);
         Self::collect(child, current, results);
         current.pop(); // backtrack
     }
 }
}

The `entry().or_insert_with()` pattern is idiomatic Rust for "insert-if-absent then get reference." It avoids a double lookup compared to `get` + `insert`. The recursive `collect` uses a mutable `current` string with `push`/`pop` — the same backtracking pattern as tree enumeration. `Box<TrieNode>` is required because the recursive type would otherwise be infinite size. The compiler cannot determine `TrieNode`'s stack size without the indirection.

What This Unlocks

Search engine autocomplete: index all query suggestions in a trie; retrieve completions in O(prefix length + results).
IP routing (longest prefix match): routers use binary tries to find the most specific matching route for a destination IP.
Spell checker suggestions: walk the trie with edit-distance tolerance to find near-matches for misspelled words.

Key Differences

Concept	OCaml	Rust
Recursive node type	`type node = { ... }` (size known by GC)	`Box<TrieNode>` — explicit indirection for recursive type
Child map	`Hashtbl.t` or `Map.Make(Char)`	`HashMap<char, Box<TrieNode>>`
Insert idiom	Pattern match + `Hashtbl.replace`	`entry().or_insert_with()` — single lookup
Autocomplete DFS	Recursive with accumulator	Mutable `current: &mut String` + `push`/`pop` backtrack
Ownership of subtree	GC-managed	`Box` owns child; parent owns `Box` via `HashMap` value

/// Trie (Prefix Tree) for autocomplete and prefix search.
///
/// Each node holds a HashMap of children and an `is_end` flag.
/// Insert and lookup are O(k) where k is the key length.

use std::collections::HashMap;

#[derive(Default, Debug)]
struct TrieNode {
    children: HashMap<char, Box<TrieNode>>,
    is_end: bool,
}

#[derive(Default, Debug)]
struct Trie {
    root: TrieNode,
}

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

    /// Insert a word into the trie.
    fn insert(&mut self, word: &str) {
        let mut node = &mut self.root;
        for ch in word.chars() {
            node = node.children.entry(ch).or_insert_with(|| Box::new(TrieNode::default()));
        }
        node.is_end = true;
    }

    /// Return true if the exact word exists.
    fn search(&self, word: &str) -> bool {
        let mut node = &self.root;
        for ch in word.chars() {
            match node.children.get(&ch) {
                Some(child) => node = child,
                None => return false,
            }
        }
        node.is_end
    }

    /// Return true if any word starts with this prefix.
    fn starts_with(&self, prefix: &str) -> bool {
        let mut node = &self.root;
        for ch in prefix.chars() {
            match node.children.get(&ch) {
                Some(child) => node = child,
                None => return false,
            }
        }
        true
    }

    /// Return all words that start with the given prefix.
    fn autocomplete(&self, prefix: &str) -> Vec<String> {
        // Navigate to the prefix node
        let mut node = &self.root;
        for ch in prefix.chars() {
            match node.children.get(&ch) {
                Some(child) => node = child,
                None => return vec![],
            }
        }
        // Collect all words in this subtrie
        let mut results = Vec::new();
        Self::collect(node, &mut prefix.to_string(), &mut results);
        results.sort();
        results
    }

    /// DFS: collect all complete words in the subtrie.
    fn collect(node: &TrieNode, current: &mut String, results: &mut Vec<String>) {
        if node.is_end {
            results.push(current.clone());
        }
        // Sort children for deterministic output
        let mut chars: Vec<char> = node.children.keys().copied().collect();
        chars.sort();
        for ch in chars {
            current.push(ch);
            Self::collect(&node.children[&ch], current, results);
            current.pop();
        }
    }
}

fn main() {
    let mut trie = Trie::new();
    let words = ["apple", "app", "application", "apply", "apt", "bat", "ball", "band"];
    for w in &words {
        trie.insert(w);
    }

    println!("autocomplete('app'): {:?}", trie.autocomplete("app"));
    println!("autocomplete('ba'):  {:?}", trie.autocomplete("ba"));
    println!("autocomplete('apt'): {:?}", trie.autocomplete("apt"));
    println!("autocomplete('xyz'): {:?}", trie.autocomplete("xyz"));
    println!("search('app'):       {}", trie.search("app"));
    println!("search('ap'):        {}", trie.search("ap"));
    println!("starts_with('ap'):   {}", trie.starts_with("ap"));
}

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

    fn make_trie() -> Trie {
        let mut t = Trie::new();
        for w in &["apple", "app", "application", "apply", "apt", "bat", "ball", "band"] {
            t.insert(w);
        }
        t
    }

    #[test]
    fn test_search_exact() {
        let t = make_trie();
        assert!(t.search("app"));
        assert!(t.search("apple"));
        assert!(t.search("apt"));
        assert!(!t.search("ap"));
        assert!(!t.search("xyz"));
    }

    #[test]
    fn test_starts_with() {
        let t = make_trie();
        assert!(t.starts_with("ap"));
        assert!(t.starts_with("ba"));
        assert!(!t.starts_with("xyz"));
    }

    #[test]
    fn test_autocomplete_app() {
        let t = make_trie();
        let mut r = t.autocomplete("app");
        r.sort();
        assert_eq!(r, vec!["app", "apple", "application", "apply"]);
    }

    #[test]
    fn test_autocomplete_ba() {
        let t = make_trie();
        assert_eq!(t.autocomplete("ba"), vec!["ball", "band", "bat"]);
    }

    #[test]
    fn test_autocomplete_missing() {
        let t = make_trie();
        assert_eq!(t.autocomplete("xyz"), Vec::<String>::new());
    }

    #[test]
    fn test_autocomplete_all() {
        let t = make_trie();
        let r = t.autocomplete("");
        assert_eq!(r.len(), 8);
    }

    #[test]
    fn test_insert_duplicate() {
        let mut t = Trie::new();
        t.insert("rust");
        t.insert("rust");
        assert_eq!(t.autocomplete("rust"), vec!["rust"]);
    }
}

(* Trie in OCaml — functional-style with records and Map *)

module CharMap = Map.Make(Char)

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

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

(* Insert a word into the trie — pure functional (returns new trie) *)
let rec insert (word : string) (pos : int) (t : trie) : trie =
  if pos = String.length word then
    { t with is_end = true }
  else
    let c = word.[pos] in
    let child =
      match CharMap.find_opt c t.children with
      | None -> empty_trie
      | Some ct -> ct
    in
    let child' = insert word (pos + 1) child in
    { t with children = CharMap.add c child' t.children }

let insert_word word t = insert word 0 t

(* Find the node for a given prefix, if it exists *)
let rec find_prefix (prefix : string) (pos : int) (t : trie) : trie option =
  if pos = String.length prefix then Some t
  else
    let c = prefix.[pos] in
    match CharMap.find_opt c t.children with
    | None -> None
    | Some child -> find_prefix prefix (pos + 1) child

(* Collect all words in a subtrie, prepending the given prefix *)
let rec collect_words (prefix : string) (t : trie) : string list =
  let from_children =
    CharMap.fold
      (fun c child acc ->
        (collect_words (prefix ^ String.make 1 c) child) @ acc)
      t.children []
  in
  if t.is_end then prefix :: from_children
  else from_children

(* Autocomplete: return all words with the given prefix *)
let autocomplete (prefix : string) (t : trie) : string list =
  match find_prefix prefix 0 t with
  | None -> []
  | Some subtrie -> List.sort String.compare (collect_words prefix subtrie)

let () =
  let words = ["apple"; "app"; "application"; "apply"; "apt"; "bat"; "ball"; "band"] in
  let trie = List.fold_right insert_word words empty_trie in

  let test prefix =
    let results = autocomplete prefix trie in
    Printf.printf "autocomplete(%S): [%s]\n" prefix (String.concat ", " results)
  in

  test "app";
  test "ba";
  test "apt";
  test "xyz";
  test ""