very_simple_trie/
lib.rs

1use std::collections::HashMap;
2
3/// A node in a Trie data structure.
4///
5/// Each node stores a map of child nodes (`children`) indexed by characters,
6/// and a boolean flag (`is_end_of_word`) indicating whether the node represents the end of a valid word.
7#[derive(Debug)]
8pub struct TrieNode {
9    children: HashMap<char, Box<TrieNode>>,
10    is_end_of_word: bool,
11}
12
13/// A Trie (prefix tree) data structure for efficient string prefix matching.
14///
15/// A Trie stores a set of strings in a tree structure, where each node represents a character in a string,
16/// and paths from the root to nodes marked as `is_end_of_word` represent complete words.
17#[derive(Debug)]
18pub struct Trie {
19    root: TrieNode,
20}
21
22impl TrieNode {
23    /// Creates a new `TrieNode` with an empty `children` map and `is_end_of_word` set to `false`.
24    pub fn new() -> Self {
25        Self {
26            children: HashMap::new(),
27            is_end_of_word: false,
28        }
29    }
30}
31
32impl Trie {
33    /// Creates a new, empty `Trie` with a root node.
34    pub fn new() -> Self {
35        Self {
36            root: TrieNode::new(),
37        }
38    }
39
40    /// Inserts a word into the `Trie`.
41    ///
42    /// This method iterates over the characters of the word, creating new nodes or following existing ones as needed.
43    /// The `is_end_of_word` flag is set to `true` on the last node to mark the end of the inserted word.
44    pub fn insert(&mut self, word: &str) {
45        let mut node = &mut self.root;
46        for char in word.chars() {
47            node = node.children.entry(char).or_insert(Box::new(TrieNode::new()));
48        }
49        node.is_end_of_word = true;
50    }
51
52    /// A private helper method for searching for a word or prefix in the Trie.
53    ///
54    /// Returns `true` if the word or prefix is found, `false` otherwise.
55    fn search(&mut self, word: &str, prefix: bool) -> bool {
56        let mut node = &self.root;
57        for char in word.chars() {
58            if let Some(next_node) = node.children.get(&char) {
59                node = next_node;
60            } else {
61                return false;
62            }
63        }
64        if prefix {
65            return true; // If searching for a prefix, any match is sufficient
66        } else {
67            node.is_end_of_word // If searching for a full word, check if we're at the end of a word
68        }
69    }
70
71    /// Searches for a prefix in the `Trie`.
72    ///
73    /// Returns `true` if the prefix exists in the `Trie`, `false` otherwise.
74    pub fn search_prefix(&mut self, word: &str) -> bool {
75        self.search(word, true)
76    }
77
78    /// Searches for a full word in the `Trie`.
79    ///
80    /// Returns `true` if the exact word exists in the `Trie`, `false` otherwise.
81    pub fn search_full_world(&mut self, word: &str) -> bool {
82        self.search(word, false)
83    }
84}
85
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90
91    #[test]
92    fn it_works() {
93        let mut trie = Trie::new();
94        trie.insert("abcde");
95        assert_eq!(trie.search_prefix("abc"), true);
96        assert_eq!(trie.search_full_world("abcde"), true);
97        assert_eq!(trie.search_full_world("abc"), false);
98    }
99}