Skip to main content

easy_trie/
trie_node.rs

1use std::collections::HashMap;
2
3#[derive(Default, Debug)]
4pub struct TrieNode {
5    pub(crate) is_end_of_word: bool,
6    pub(crate) children: HashMap<char, TrieNode>,
7}
8
9impl TrieNode {
10    /**
11     * Calculates the length of the current node.
12     *
13     * Defined as the sum of the length of all children nodes and of the current node.
14     */
15    pub(crate) fn len(&self) -> usize {
16        let mut sum = self.children.len();
17        for a in self.children.values() {
18            sum += a.len();
19        }
20        sum
21    }
22
23    /**
24     * Get all possible words for a given node.
25     */
26    pub(crate) fn get_suggestions(&self, current: &str) -> Vec<String> {
27        let mut words: Vec<String> = Vec::new();
28        if self.is_end_of_word {
29            words.push(current.to_string());
30            eprintln!("Inserting {current}")
31        }
32        for (c, node) in self.children.iter() {
33            words.append(
34                node.get_suggestions(format!("{current}{c}").as_str())
35                    .as_mut(),
36            );
37        }
38        eprintln!("returning {words:#?}");
39        words
40    }
41}
42
43#[cfg(test)]
44mod tests {
45    use super::*;
46
47    #[test]
48    fn it_should_create_a_new_node() {
49        let result = TrieNode::default();
50        assert!(!result.is_end_of_word);
51        assert_eq!(result.children.len(), 0);
52    }
53
54    #[test]
55    fn it_should_have_correct_length() {
56        let mut node = TrieNode::default();
57        assert_eq!(node.len(), 0);
58        let mut child = TrieNode::default();
59        child.is_end_of_word = true;
60        node.children.insert('a', child);
61        assert_eq!(node.len(), 1);
62    }
63
64    #[test]
65    fn it_should_get_suggestions() {
66        let mut node = TrieNode::default();
67        let mut child = TrieNode::default();
68        child.is_end_of_word = true;
69        node.children.insert('a', child);
70        let words = node.get_suggestions("");
71        assert_eq!(words, vec!["a"]);
72    }
73
74    #[test]
75    fn it_should_get_suggestions_independantly_from_prefix() {
76        let mut node = TrieNode::default();
77        let mut child = TrieNode::default();
78        child.is_end_of_word = true;
79        node.children.insert('a', child);
80        let prefixes = ["some", "prefix"];
81        for prefix in prefixes.iter() {
82            let words = node.get_suggestions(prefix);
83            assert_eq!(words, vec![format!("{prefix}a")]);
84        }
85    }
86
87    #[test]
88    fn it_should_get_suggestions_with_multiple_children() {
89        let mut node = TrieNode::default();
90        let mut child = TrieNode::default();
91        child.is_end_of_word = true;
92        node.children.insert('a', child);
93        let mut child = TrieNode::default();
94        child.is_end_of_word = true;
95        node.children.insert('b', child);
96        let words = node.get_suggestions("");
97        assert!(words.contains(&"a".to_string()));
98        assert!(words.contains(&"b".to_string()));
99    }
100}