forked from neetcode-gh/leetcode
-
Notifications
You must be signed in to change notification settings - Fork 0
/
0208-implement-trie-prefix-tree.rs
56 lines (44 loc) · 1.21 KB
/
0208-implement-trie-prefix-tree.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
#[derive(Default)]
struct Trie {
is_word_end: bool,
children: [Option<Box<Trie>>; 26],
}
impl Trie {
fn new() -> Self {
Self{ ..Default::default() }
}
fn insert(&mut self, word: String) {
let mut trie = self;
for c in word.chars(){
trie = trie.children[char_index(c)]
.get_or_insert(Box::new(Trie::new()))
.as_mut();
}
trie.is_word_end = true;
}
fn search(& self, word: String) -> bool {
let mut trie = self;
for c in word.chars(){
if let Some(ref next_trie) = trie.children[char_index(c)]{
trie = next_trie.as_ref();
}else{
return false;
}
}
trie.is_word_end
}
fn starts_with(&self, prefix: String) -> bool {
let mut trie = self;
for c in prefix.chars(){
if let Some(ref next_trie) = trie.children[char_index(c)]{
trie = next_trie.as_ref();
}else{
return false;
}
}
true
}
}
pub fn char_index(c: char) -> usize{
c as usize - 'a' as usize
}