aoclib_rs/
trie.rs

1use std::collections::HashMap;
2
3/// Recursively-defined prefix trie. Typically constructed with `Trie::new()` for the root node,
4/// then `insert()`ed into many times. Querying for simple existence of a word in the `Trie` can be
5/// performed with `contains()`.
6///
7/// More advanced usage includes using `find()` to find sub-`Trie` structures to efficiently search
8/// for existince of many words with similar prefixes. See
9/// <https://github.com/glennhartmann/aoc24/blob/fa21b5787382765a598381c3a0583258b645dc86/src/days/day_19.rs>
10/// for an example.
11///
12/// Empty strings are considered to exist in every `Trie` by default.
13pub struct Trie {
14    prefix: String,
15    children: HashMap<char, Self>,
16    is_terminal: bool,
17}
18
19impl Trie {
20    /// Create a new `Trie` root node containing the empty string.
21    pub fn new() -> Self {
22        Self::with_prefix("", true /* is_terminal */)
23    }
24
25    /// Create a sub-node with an existing prefix, whether terminal or not, or can be used to
26    /// create a root node which does not contain the empty string.
27    pub fn with_prefix(s: &str, is_terminal: bool) -> Self {
28        Self {
29            prefix: s.to_owned(),
30            children: HashMap::new(),
31            is_terminal,
32        }
33    }
34
35    /// Insert a word into the `Trie`.
36    pub fn insert(&mut self, word: &str) {
37        let mut curr = self;
38        let chars: Vec<_> = word.chars().collect();
39        for (i, c) in chars.iter().enumerate() {
40            let is_terminal = i == chars.len() - 1;
41            curr = curr
42                .children
43                .entry(*c)
44                .and_modify(|t| t.is_terminal |= is_terminal)
45                .or_insert(Self::with_prefix(&word[..=i], is_terminal));
46        }
47    }
48
49    /// Find the node within the `Trie` described by `word`. Note that the returned node may not be
50    /// terminal. Can be used as any other `Trie`, for example to continue searching for many
51    /// suffixes that share the same prefix.
52    pub fn find(&self, word: &str) -> Option<&Self> {
53        let mut curr = self;
54        for c in word.chars() {
55            match curr.children.get(&c) {
56                None => return None,
57                Some(t) => curr = t,
58            };
59        }
60
61        Some(curr)
62    }
63
64    /// Whether the `Trie` contains `word`.
65    pub fn contains(&self, word: &str) -> bool {
66        let Some(node) = self.find(word) else {
67            return false;
68        };
69
70        node.is_terminal()
71    }
72
73    pub fn is_terminal(&self) -> bool {
74        self.is_terminal
75    }
76
77    pub fn prefix(&self) -> &str {
78        &self.prefix
79    }
80}
81
82impl Default for Trie {
83    fn default() -> Self {
84        Self::new()
85    }
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91
92    #[test]
93    fn test_empty() {
94        let t = Trie::new();
95        assert!(t.contains(""));
96        assert!(!t.contains("a"));
97    }
98
99    #[test]
100    fn test_typical() {
101        let mut t = Trie::new();
102        assert!(t.contains(""));
103
104        t.insert("asdf");
105        t.insert("1234");
106        t.insert("asteroid");
107        t.insert("astronaut");
108        t.insert("1235");
109
110        assert!(t.contains(""));
111        assert!(t.contains("asdf"));
112        assert!(t.contains("1234"));
113        assert!(t.contains("asteroid"));
114        assert!(t.contains("astronaut"));
115        assert!(t.contains("1235"));
116
117        assert!(!t.contains("as"));
118        t.insert("as");
119        assert!(t.contains("as"));
120        assert!(t.contains("asteroid"));
121        assert!(t.contains("astronaut"));
122        assert!(t.contains("asdf"));
123    }
124}
aoclib_rs/trie.rs

aoclib_rs/
trie.rs
