-
Notifications
You must be signed in to change notification settings - Fork 73
/
trie.rs
97 lines (82 loc) · 2.24 KB
/
trie.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
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
use std::collections::HashMap;
use std::hash::Hash;
#[derive(Debug, Default)]
struct Node<Key: Default, Type: Default> {
children: HashMap<Key, Node<Key, Type>>,
value: Option<Type>,
}
#[derive(Debug, Default)]
pub struct Trie<Key, Type>
where
Key: Default + Eq + Hash,
Type: Default,
{
root: Node<Key, Type>,
}
impl<Key, Type> Trie<Key, Type>
where
Key: Default + Eq + Hash,
Type: Default,
{
pub fn new() -> Self {
Self {
root: Node::default(),
}
}
pub fn insert(&mut self, key: impl IntoIterator<Item = Key>, value: Type)
where
Key: Eq + Hash,
{
let mut node = &mut self.root;
for c in key.into_iter() {
node = node.children.entry(c).or_default();
}
node.value = Some(value);
}
pub fn get(&self, key: impl IntoIterator<Item = Key>) -> Option<&Type>
where
Key: Eq + Hash,
{
let mut node = &self.root;
for c in key.into_iter() {
if node.children.contains_key(&c) {
node = node.children.get(&c).unwrap()
} else {
return None;
}
}
node.value.as_ref()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_insertion() {
let mut trie = Trie::new();
assert_eq!(trie.get("".chars()), None);
trie.insert("foo".chars(), 1);
trie.insert("foobar".chars(), 2);
let mut trie = Trie::new();
assert_eq!(trie.get(vec![1, 2, 3]), None);
trie.insert(vec![1, 2, 3], 1);
trie.insert(vec![3, 4, 5], 2);
}
#[test]
fn test_get() {
let mut trie = Trie::new();
trie.insert("foo".chars(), 1);
trie.insert("foobar".chars(), 2);
trie.insert("bar".chars(), 3);
trie.insert("baz".chars(), 4);
assert_eq!(trie.get("foo".chars()), Some(&1));
assert_eq!(trie.get("food".chars()), None);
let mut trie = Trie::new();
trie.insert(vec![1, 2, 3, 4], 1);
trie.insert(vec![42], 2);
trie.insert(vec![42, 6, 1000], 3);
trie.insert(vec![1, 2, 4, 16, 32], 4);
assert_eq!(trie.get(vec![42, 6, 1000]), Some(&3));
assert_eq!(trie.get(vec![43, 44, 45]), None);
}
}