trie/
lib.rs

1//! Trie (prefix tree) for storing and querying strings.
2//!
3//! Supports three operations, each `O(L)` where `L` is the length of the input:
4//! - [`Trie::insert`] - add a word to the trie
5//! - [`Trie::search`] - whether a word was previously inserted
6//! - [`Trie::starts_with`] - whether any inserted word has the given prefix
7//!
8//! Children are stored in a `HashMap<char, TrieNode>` for `O(1)` descent.
9
10use std::collections::HashMap;
11
12#[derive(Default)]
13struct TrieNode {
14    children: HashMap<char, TrieNode>,
15    is_end_of_word: bool,
16}
17
18/// Trie data structure.
19#[derive(Default)]
20pub struct Trie {
21    root: TrieNode,
22}
23
24impl Trie {
25    /// Construct an empty trie.
26    pub fn new() -> Self {
27        Default::default()
28    }
29
30    /// Insert `word` into the trie.
31    pub fn insert(&mut self, word: &str) {
32        let mut curr = &mut self.root;
33        for ch in word.chars() {
34            curr = curr.children.entry(ch).or_default();
35        }
36        curr.is_end_of_word = true;
37    }
38
39    /// Whether `word` was previously inserted (and not merely a prefix).
40    pub fn search(&self, word: &str) -> bool {
41        let mut curr = &self.root;
42        for ch in word.chars() {
43            match curr.children.get(&ch) {
44                Some(node) => curr = node,
45                None => return false,
46            }
47        }
48        curr.is_end_of_word
49    }
50
51    /// Whether any inserted word has `prefix`.
52    pub fn starts_with(&self, prefix: &str) -> bool {
53        let mut curr = &self.root;
54        for ch in prefix.chars() {
55            match curr.children.get(&ch) {
56                Some(node) => curr = node,
57                None => return false,
58            }
59        }
60        true
61    }
62}
trie/lib.rs

trie/
lib.rs
