competitive_programming_lib/DataStructures/
trie.rs

1#[derive(Default)]
2pub struct TrieNode {
3    children: std::collections::HashMap<char, TrieNode>,
4    is_end_of_word: bool,
5}
6
7#[derive(Default)]
8pub struct Trie {
9    root: TrieNode,
10}
11
12impl Trie {
13    pub fn new() -> Self {
14        Trie {
15            root: TrieNode::default(),
16        }
17    }
18
19    pub fn insert(&mut self, word: &str) {
20        let mut current_node = &mut self.root;
21        for c in word.chars() {
22            current_node = current_node
23                .children
24                .entry(c)
25                .or_insert_with(TrieNode::default);
26        }
27        current_node.is_end_of_word = true;
28    }
29
30    pub fn search(&self, word: &str) -> bool {
31        let mut current_node = &self.root;
32        for c in word.chars() {
33            if let Some(node) = current_node.children.get(&c) {
34                current_node = node;
35            } else {
36                return false;
37            }
38        }
39        current_node.is_end_of_word
40    }
41
42    pub fn starts_with(&self, prefix: &str) -> bool {
43        let mut current_node = &self.root;
44        for c in prefix.chars() {
45            if let Some(node) = current_node.children.get(&c) {
46                current_node = node;
47            } else {
48                return false;
49            }
50        }
51        true
52    }
53}