962: Trie Map

Difficulty: Intermediate Category: Data Structures / Trees Concept: Prefix tree with per-node optional values, supporting insert, search, and prefix queries Key Insight: OCaml uses `CharMap` (balanced BST over chars) for children; Rust uses `HashMap<char, TrieNode<V>>` — both achieve prefix tree semantics, but Rust's `entry().or_insert_with()` idiom elegantly handles the "get or create child" pattern

// 962: Trie Map
// OCaml: mutable record with CharMap; Rust: struct with HashMap children

use std::collections::HashMap;

#[derive(Debug)]
pub struct TrieNode<V> {
    value: Option<V>,
    children: HashMap<char, TrieNode<V>>,
}

impl<V> Default for TrieNode<V> {
    fn default() -> Self {
        TrieNode {
            value: None,
            children: HashMap::new(),
        }
    }
}

pub struct Trie<V> {
    root: TrieNode<V>,
}

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

    pub fn insert(&mut self, key: &str, value: V) {
        let mut node = &mut self.root;
        for c in key.chars() {
            node = node.children.entry(c).or_default();
        }
        node.value = Some(value);
    }

    pub fn search(&self, key: &str) -> Option<&V> {
        let mut node = &self.root;
        for c in key.chars() {
            match node.children.get(&c) {
                Some(child) => node = child,
                None => return None,
            }
        }
        node.value.as_ref()
    }

    pub fn starts_with(&self, prefix: &str) -> bool {
        let mut node = &self.root;
        for c in prefix.chars() {
            match node.children.get(&c) {
                Some(child) => node = child,
                None => return false,
            }
        }
        true
    }

    // Approach 2: Collect all keys with a given prefix
    pub fn keys_with_prefix(&self, prefix: &str) -> Vec<String> {
        let mut node = &self.root;
        for c in prefix.chars() {
            match node.children.get(&c) {
                Some(child) => node = child,
                None => return vec![],
            }
        }
        let mut results = vec![];
        Self::collect_keys(node, &mut prefix.to_string(), &mut results);
        results
    }

    fn collect_keys(node: &TrieNode<V>, prefix: &mut String, results: &mut Vec<String>) {
        if node.value.is_some() {
            results.push(prefix.clone());
        }
        let mut children: Vec<(&char, &TrieNode<V>)> = node.children.iter().collect();
        children.sort_by_key(|(c, _)| *c);
        for (c, child) in children {
            prefix.push(*c);
            Self::collect_keys(child, prefix, results);
            prefix.pop();
        }
    }
}

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


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

    fn make_trie() -> Trie<i32> {
        let mut t = Trie::new();
        t.insert("apple", 1);
        t.insert("app", 2);
        t.insert("application", 3);
        t.insert("banana", 4);
        t.insert("band", 5);
        t
    }

    #[test]
    fn test_search_found() {
        let t = make_trie();
        assert_eq!(t.search("apple"), Some(&1));
        assert_eq!(t.search("app"), Some(&2));
        assert_eq!(t.search("application"), Some(&3));
        assert_eq!(t.search("banana"), Some(&4));
        assert_eq!(t.search("band"), Some(&5));
    }

    #[test]
    fn test_search_not_found() {
        let t = make_trie();
        assert_eq!(t.search("ap"), None);
        assert_eq!(t.search("apricot"), None);
        assert_eq!(t.search(""), None);
        assert_eq!(t.search("xyz"), None);
    }

    #[test]
    fn test_starts_with() {
        let t = make_trie();
        assert!(t.starts_with("app"));
        assert!(t.starts_with("ban"));
        assert!(t.starts_with("apple"));
        assert!(!t.starts_with("xyz"));
        assert!(!t.starts_with("apricot"));
    }

    #[test]
    fn test_update() {
        let mut t = make_trie();
        t.insert("apple", 99);
        assert_eq!(t.search("apple"), Some(&99));
    }

    #[test]
    fn test_prefix_keys() {
        let t = make_trie();
        let mut keys = t.keys_with_prefix("app");
        keys.sort();
        assert_eq!(keys, vec!["app", "apple", "application"]);

        let ban_keys = t.keys_with_prefix("ban");
        assert_eq!(ban_keys.len(), 2);
        assert!(ban_keys.contains(&"banana".to_string()));
        assert!(ban_keys.contains(&"band".to_string()));
    }
}

(* 962: Trie Map *)
(* Prefix tree for string keys with associated values *)

module CharMap = Map.Make(Char)

type 'a trie = {
  mutable value: 'a option;
  mutable children: 'a trie CharMap.t;
}

let create_node () = { value = None; children = CharMap.empty }

let insert trie key value =
  let node = ref trie in
  String.iter (fun c ->
    let child =
      match CharMap.find_opt c (!node).children with
      | Some n -> n
      | None ->
        let n = create_node () in
        (!node).children <- CharMap.add c n (!node).children;
        n
    in
    node := child
  ) key;
  (!node).value <- Some value

let search trie key =
  let node = ref (Some trie) in
  let i = ref 0 in
  let n = String.length key in
  while !i < n && !node <> None do
    (match !node with
     | Some nd ->
       node := CharMap.find_opt key.[!i] nd.children
     | None -> ());
    incr i
  done;
  match !node with
  | Some nd -> nd.value
  | None -> None

let starts_with trie prefix =
  let node = ref (Some trie) in
  let i = ref 0 in
  let n = String.length prefix in
  while !i < n && !node <> None do
    (match !node with
     | Some nd ->
       node := CharMap.find_opt prefix.[!i] nd.children
     | None -> ());
    incr i
  done;
  !node <> None

let () =
  let t = create_node () in
  insert t "apple" 1;
  insert t "app" 2;
  insert t "application" 3;
  insert t "banana" 4;
  insert t "band" 5;

  assert (search t "apple" = Some 1);
  assert (search t "app" = Some 2);
  assert (search t "application" = Some 3);
  assert (search t "banana" = Some 4);
  assert (search t "band" = Some 5);
  assert (search t "ap" = None);      (* prefix, not full word *)
  assert (search t "apricot" = None);
  assert (search t "" = None);

  assert (starts_with t "app");
  assert (starts_with t "ban");
  assert (starts_with t "apple");
  assert (not (starts_with t "xyz"));
  assert (not (starts_with t "apricot"));

  (* Update existing key *)
  insert t "apple" 99;
  assert (search t "apple" = Some 99);

  Printf.printf "✓ All tests passed\n"

📊 Detailed Comparison

Trie Map — Comparison

Core Insight

A trie stores strings by sharing common prefixes — each node is a character + optional value + map of children. The algorithm is identical in both languages. OCaml uses a functional `Map.Make(Char)` (immutable BST) or mutable fields; Rust uses `HashMap<char, TrieNode<V>>` with `entry().or_insert_with()` for clean "get or create" logic.

OCaml Approach

`module CharMap = Map.Make(Char)` — ordered map over chars
Mutable fields (`mutable value`, `mutable children`) on a record
Imperative traversal with `ref` node pointer
`String.iter (fun c -> ...)` to walk characters
`CharMap.find_opt` / `CharMap.add` for children

Rust Approach

`HashMap<char, TrieNode<V>>` for children (O(1) vs O(log n) for BST)
`node.children.entry(c).or_insert_with(TrieNode::default)` — get-or-create idiom
`for c in key.chars()` — Unicode-safe char iteration
`let mut node = &mut self.root` — mutable reference traversal
Generic over `V` with `Default` bound for empty nodes
`collect_keys` recursive DFS for prefix enumeration

Comparison Table

Aspect	OCaml	Rust
Children map	`CharMap.t` (balanced BST)	`HashMap<char, TrieNode<V>>`
Get-or-create child	`match find_opt c children with None -> create`	`entry(c).or_insert_with(Default::default)`
Node pointer	`let node = ref trie`	`let mut node = &mut self.root`
Char iteration	`String.iter (fun c -> ...)`	`for c in key.chars()`
Value storage	`mutable value: 'a option`	`value: Option<V>`
Lookup result	`nd.value` (option)	`node.value.as_ref()`
Prefix DFS	Manual recursion	Recursive helper with `&mut String`