Prefix Search Using Trie in DSA Rust

Mental Model

A trie stores words by sharing common beginnings, so searching for words starting with a prefix is fast and easy.

Analogy: Imagine a tree of roads where each road sign is a letter. To find all streets starting with 'Ma', you just follow the 'M' road, then the 'a' road, and see all streets branching from there.

root
 ↓
 m -> a -> n -> null
      ↓
      p -> null
      ↓
      e -> null

Dry Run Walkthrough

Input: Insert words: "man", "map", "made"; then search prefix "ma"

Goal: Find all words starting with prefix "ma"

Step 1: Insert 'man' letter by letter

root -> m -> a -> n -> null

Why: Build path for 'man' so trie stores it

Step 2: Insert 'map' sharing 'm' and 'a', add 'p'

root -> m -> a -> n -> null
               ↓
               p -> null

Why: Reuse existing prefix 'ma' to save space

Step 3: Insert 'made' sharing 'm' and 'a', add 'd' and 'e'

root -> m -> a -> n -> null
               ↓
               p -> null
               ↓
               d -> e -> null

Why: Add new branch for 'made' after 'ma'

Step 4: Search prefix 'ma' by following 'm' then 'a'

At node 'a' with children n, p, d

Why: Find node representing prefix 'ma' to get all words below

Step 5: Collect all words below 'a': 'man', 'map', 'made'

Words found: man, map, made

Why: All words starting with 'ma' are under this node

Result:

root -> m -> a -> n -> null
               ↓
               p -> null
               ↓
               d -> e -> null
Words starting with 'ma': man, map, made

Annotated Code

DSA Rust

use std::collections::HashMap;

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

impl TrieNode {
    fn new() -> Self {
        TrieNode {
            children: HashMap::new(),
            is_end: false,
        }
    }
}

struct Trie {
    root: TrieNode,
}

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

    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(TrieNode::new); // advance or create child
        }
        node.is_end = true; // mark end of word
    }

    fn search_prefix_node(&self, prefix: &str) -> Option<&TrieNode> {
        let mut node = &self.root;
        for ch in prefix.chars() {
            match node.children.get(&ch) {
                Some(next_node) => node = next_node, // advance to next node
                None => return None, // prefix not found
            }
        }
        Some(node) // return node at prefix end
    }

    fn collect_words(&self, node: &TrieNode, prefix: &mut String, results: &mut Vec<String>) {
        if node.is_end {
            results.push(prefix.clone()); // found a word
        }
        for (ch, child) in &node.children {
            prefix.push(*ch); // add char
            self.collect_words(child, prefix, results); // recurse
            prefix.pop(); // backtrack
        }
    }

    fn starts_with(&self, prefix: &str) -> Vec<String> {
        let mut results = Vec::new();
        if let Some(node) = self.search_prefix_node(prefix) {
            let mut prefix_string = prefix.to_string();
            self.collect_words(node, &mut prefix_string, &mut results); // gather words
        }
        results
    }
}

fn main() {
    let mut trie = Trie::new();
    trie.insert("man");
    trie.insert("map");
    trie.insert("made");

    let prefix = "ma";
    let words = trie.starts_with(prefix);
    println!("Words starting with '{}': {}", prefix, words.join(", "));
}

node = node.children.entry(ch).or_insert_with(TrieNode::new);

advance or create child node for current character

node.is_end = true;

mark current node as end of a word

match node.children.get(&ch) { ... }

traverse children to find prefix path

results.push(prefix.clone());

add found word to results when end node reached

self.collect_words(child, prefix, results);

recursively collect all words below current node

Output

Words starting with 'ma': man, map, made

Complexity Analysis

Time: O(m + k) where m is length of prefix and k is total characters in matched words because we traverse prefix then collect all words below

Space: O(n) where n is total characters stored in trie nodes for all inserted words

vs Alternative: Compared to scanning all words linearly (O(w * l)), trie reduces prefix search to O(m + k), much faster for large word sets

Edge Cases

Empty prefix search

Returns all words stored in trie

DSA Rust

if let Some(node) = self.search_prefix_node(prefix) {

Prefix not found

Returns empty list

DSA Rust

None => return None,

Insert empty string

Marks root as end of word, allowing empty word search

DSA Rust

node.is_end = true;

When to Use This Pattern

When you need to quickly find all words starting with a prefix, use a trie because it shares common prefixes and speeds up search.

Common Mistakes

Mistake: Not marking the end of a word, causing prefix search to return incomplete results
Fix: Set is_end = true at the last node of inserted word

Mistake: Not handling prefix not found case, leading to errors
Fix: Return None or empty list when prefix path is missing

Summary

Stores words in a tree structure sharing prefixes for fast prefix search.

Use when you want to find all words starting with a given prefix efficiently.

The key is to follow the prefix path in the trie, then collect all words below that node.

Practice

1. What is the main purpose of using a Trie data structure in prefix search?

easy

A. To efficiently find all words that start with a given prefix

B. To sort words alphabetically

C. To store numbers in sorted order

D. To perform arithmetic operations quickly

Solution

1. Step 1: Understand Trie usage

   A Trie stores strings in a tree-like structure where each node represents a character, making prefix searches efficient.

2. Step 2: Identify prefix search goal

   Prefix search requires finding all words starting with a given prefix, which Trie supports directly by traversing nodes.

3. Final Answer:

   To efficiently find all words that start with a given prefix -> Option A

4. Quick Check:

   Trie = prefix search efficiency [OK]

Hint: Trie nodes represent characters for fast prefix lookup [OK]

Common Mistakes:

• Confusing Trie with sorting data structure
• Thinking Trie is for numeric operations
• Assuming Trie stores full words only at root

2. Which of the following Rust code snippets correctly defines a Trie node struct for prefix search?

easy

A. struct TrieNode { children: HashMap, is_end: bool }

B. struct TrieNode { children: Vec<char>, is_end: bool }

C. struct TrieNode { children: String, is_end: bool }

D. struct TrieNode { children: i32, is_end: bool }

Solution

1. Step 1: Check children type for TrieNode

   Children must map characters to child nodes, so HashMap is correct for recursive structure without size issues.

2. Step 2: Validate other options

   Vec<char> or String cannot map characters to nodes; i32 is invalid for children.

3. Final Answer:

   struct TrieNode { children: HashMap, is_end: bool } -> Option A

4. Quick Check:

   Children map chars to nodes = struct TrieNode { children: HashMap, is_end: bool } [OK]

Hint: Use HashMap for children in Rust Trie [OK]

Common Mistakes:

• Using Vec<char> instead of mapping to nodes
• Using String for children instead of node references
• Using primitive types like i32 for children

3. Given the following Rust code snippet for inserting words into a Trie and searching for prefix "ca":

let mut trie = Trie::new();
trie.insert("cat");
trie.insert("car");
trie.insert("dog");
let result = trie.starts_with("ca");
println!("{:?}", result);

What will be printed?

medium

A. ["cat", "car", "dog"]

B. ["cat", "car"]

C. ["dog"]

D. []

Solution

1. Step 1: Insert words into Trie

   Words "cat", "car", and "dog" are inserted. "cat" and "car" share prefix "ca".

2. Step 2: Search for prefix "ca"

   Searching prefix "ca" returns all words starting with "ca", which are "cat" and "car".

3. Final Answer:

   ["cat", "car"] -> Option B

4. Quick Check:

   Prefix "ca" matches cat, car [OK]

Hint: Prefix search returns words starting exactly with given prefix [OK]

Common Mistakes:

• Including words without the prefix
• Returning empty list if prefix exists
• Confusing prefix search with full word search

4. Consider this Rust code snippet for prefix search in a Trie:

fn starts_with(&self, prefix: &str) -> bool {
    let mut node = &self.root;
    for ch in prefix.chars() {
        if !node.children.contains_key(&ch) {
            return false;
        }
        node = node.children.get(&ch).unwrap();
    }
    true
}

What is the main issue with this code snippet?

medium

A. Cannot assign &TrieNode to node because of lifetime mismatch

B. Missing mutable keyword for node variable

C. Using unwrap() without checking Option

D. Incorrect return type for starts_with function

Solution

1. Step 1: Analyze function purpose

   This implements prefix existence check (returns bool).

2. Step 2: Identify mismatch for prefix search

   Prefix search requires returning all words with the prefix (e.g. Vec<String>), not just bool.

3. Final Answer:

   Incorrect return type for starts_with function -> Option D

4. Quick Check:

   Prefix search needs Vec<String> not bool [OK]

Hint: Verify return type matches prefix search requirements [OK]

Common Mistakes:

• Confusing prefix existence (bool) with full prefix search (list)
• Assuming unwrap() causes runtime error here
• Forgetting Rust borrow checker rules

5. You want to extend a Trie in Rust to support prefix search that returns all words starting with a prefix, but also limit results to a maximum of 2 words. Which approach is best?

hard

A. Use a HashSet to store all words and filter after search

B. Store only 2 words per node during insertion

C. Modify DFS traversal to stop collecting words after 2 matches

D. Return only the first word found with the prefix

Solution

1. Step 1: Understand prefix search with limit

   We want to find all words starting with prefix but stop after 2 matches to save time.

2. Step 2: Modify DFS traversal

   During DFS from prefix node, keep count of collected words and stop recursion once 2 words are found.

3. Step 3: Evaluate other options

   Storing only 2 words per node loses data; filtering after search is inefficient; returning one word is incomplete.

4. Final Answer:

   Modify DFS traversal to stop collecting words after 2 matches -> Option C

5. Quick Check:

   Limit results by stopping DFS early [OK]

Hint: Stop DFS early to limit prefix search results [OK]

Common Mistakes:

• Limiting stored words per node loses data
• Filtering after full search wastes time
• Returning only one word ignores requirement