collect_me/tree/
binary_tree.rs

1/// A binary tree containing key-value pairs where the keys can be ordered.
2///
3/// It should be noted that for most applications, a [`std::collections::HashMap`] will offer
4/// superior performance to that of a binary tree, since each node in the tree requires a heap
5/// allocation (apart from the root). Hash maps also provided amortized-constant lookup times where
6/// a binary tree gives `O(log(n))`.
7///
8/// For efficiency, the tree maintains a count of the number of elements inserted so that the
9/// `len` and `is_empty` methods are constant-time complexity.
10///
11/// # Examples
12///
13/// This example shows how the binary tree functions much like a [`std::collections::HashMap`], but
14/// gives `O(log(n))` lookup time for keys that are of an ordinal type.
15/// ```
16/// use collect_me::tree::binary_tree::BinaryTree;
17///
18/// let mut tree = BinaryTree::new();
19/// tree.insert(0, "John");
20/// tree.insert(42, "Neo");
21/// tree.insert(2, "Alice");
22///
23/// assert_eq!(tree.get(&0), Some(&"John"));
24/// assert_eq!(tree.get(&42), Some(&"Neo"));
25/// assert_eq!(tree.get(&2), Some(&"Alice"));
26/// ```
27#[derive(Debug, Clone, Default, PartialEq, Eq)]
28pub struct BinaryTree<K, V> {
29    root: Option<BinaryTreeNode<K, V>>,
30    len: usize,
31}
32
33type NodeChild<K, V> = Option<Box<BinaryTreeNode<K, V>>>;
34
35#[doc(hidden)]
36#[derive(Debug, Clone, Default, PartialEq, Eq)]
37struct BinaryTreeNode<K, V> {
38    key: K,
39    value: V,
40    children: (NodeChild<K, V>, NodeChild<K, V>),
41}
42
43impl<K, V> BinaryTree<K, V> {
44    /// Constructs an empty tree
45    pub fn new() -> Self {
46        Self { root: None, len: 0 }
47    }
48}
49
50impl<K, V> BinaryTree<K, V>
51where
52    K: PartialOrd + Eq,
53{
54    /// Inserts a key-value pair into the [`BinaryTree`].
55    ///
56    /// Returns [`None`] if the key did not exist, otherwise updates
57    /// the value and returns [`Some`] with the old value.
58    ///
59    /// # Note
60    ///
61    /// Like with [`std::collections::HashMap`] the key does not get updated.
62    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
63        let result = if let Some(ref mut root) = self.root {
64            root.insert(key, value)
65        } else {
66            self.root = Some(BinaryTreeNode {
67                key,
68                value,
69                children: (None, None),
70            });
71            None
72        };
73
74        if result.is_none() {
75            self.len += 1;
76        }
77
78        result
79    }
80
81    /// Returns a reference to the value corresponding to the key.
82    pub fn get<Q>(&self, key: &Q) -> Option<&V>
83    where
84        K: std::borrow::Borrow<Q>,
85        Q: PartialOrd + Eq,
86    {
87        self.root.as_ref().and_then(|root| root.get(key))
88    }
89
90    /// Returns a mutable reference to the value corresponding to the key.
91    pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
92    where
93        K: std::borrow::Borrow<Q>,
94        Q: PartialOrd + Eq,
95    {
96        self.root.as_mut().and_then(|root| root.get_mut(key))
97    }
98
99    /// Returns the number of elements in the tree with constant-time complexity.
100    pub fn len(&self) -> usize {
101        self.len
102    }
103
104    /// Returns `true` if the tree is empty.
105    pub fn is_empty(&self) -> bool {
106        self.len == 0
107    }
108}
109
110impl<K, V> BinaryTreeNode<K, V>
111where
112    K: PartialOrd + Eq,
113{
114    fn insert(&mut self, key: K, value: V) -> Option<V> {
115        if key < self.key {
116            if let Some(ref mut child) = self.children.0 {
117                child.insert(key, value)
118            } else {
119                self.children.0 = Some(Box::new(BinaryTreeNode {
120                    key,
121                    value,
122                    children: (None, None),
123                }));
124                None
125            }
126        } else if key > self.key {
127            if let Some(ref mut child) = self.children.1 {
128                child.insert(key, value)
129            } else {
130                self.children.1 = Some(Box::new(BinaryTreeNode {
131                    key,
132                    value,
133                    children: (None, None),
134                }));
135                None
136            }
137        } else {
138            Some(std::mem::replace(&mut self.value, value))
139        }
140    }
141
142    fn get<Q>(&self, key: &Q) -> Option<&V>
143    where
144        K: std::borrow::Borrow<Q>,
145        Q: PartialOrd + Eq,
146    {
147        if *key == *self.key.borrow() {
148            Some(&self.value)
149        } else if *key < *self.key.borrow() {
150            self.children.0.as_ref().and_then(|child| child.get(key))
151        } else {
152            self.children.1.as_ref().and_then(|child| child.get(key))
153        }
154    }
155
156    fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
157    where
158        K: std::borrow::Borrow<Q>,
159        Q: PartialOrd + Eq,
160    {
161        if *key == *self.key.borrow() {
162            Some(&mut self.value)
163        } else if *key < *self.key.borrow() {
164            self.children
165                .0
166                .as_mut()
167                .and_then(|child| child.get_mut(key))
168        } else {
169            self.children
170                .1
171                .as_mut()
172                .and_then(|child| child.get_mut(key))
173        }
174    }
175}
176
177impl<K, V> std::ops::Index<&K> for BinaryTree<K, V>
178where
179    K: PartialOrd + Eq,
180{
181    type Output = V;
182
183    /// Returns a reference to the value corresponding to the supplied key.
184    ///
185    /// # Panics
186    ///
187    /// Panics if the key is not present in the binary tree.
188    fn index(&self, index: &K) -> &Self::Output {
189        self.get(index)
190            .expect("Key is not present in the binary tree")
191    }
192}
193
194#[cfg(test)]
195mod tests {
196    use super::*;
197
198    #[test]
199    fn empty() {
200        let tree: BinaryTree<i32, i32> = BinaryTree::new();
201        assert_eq!(tree.get(&0), None);
202        assert_eq!(tree.len(), 0);
203        assert!(tree.is_empty());
204    }
205
206    #[test]
207    fn one_key() {
208        let mut tree = BinaryTree::new();
209        tree.insert(3, "Hello");
210        assert_eq!(tree.get(&3), Some(&"Hello"));
211        assert_eq!(tree.insert(3, "World"), Some("Hello"));
212        assert_eq!(tree.get(&3), Some(&"World"));
213    }
214
215    #[test]
216    fn three_keys() {
217        let mut tree = BinaryTree::new();
218        tree.insert(1, 'A');
219        tree.insert(0, 'B');
220        tree.insert(2, 'C');
221
222        assert_eq!(tree.get(&1), Some(&'A'));
223        assert_eq!(tree.get(&0), Some(&'B'));
224        assert_eq!(tree.get(&2), Some(&'C'));
225        assert_eq!(tree.get(&-1), None);
226    }
227
228    #[test]
229    fn get_mut() {
230        let mut tree = BinaryTree::new();
231        tree.insert(1, 'A');
232        tree.insert(2, 'B');
233        tree.insert(0, 'C');
234
235        assert_eq!(tree.get(&0), Some(&'C'));
236
237        let val = tree.get_mut(&1).expect("Failed to mutably reference value");
238        *val = 'X';
239
240        let val = tree.get_mut(&2).expect("Failed to mutably reference value");
241        *val = 'Y';
242
243        let val = tree.get_mut(&0).expect("Failed to mutably reference value");
244        *val = 'Z';
245
246        assert_eq!(tree.get(&1), Some(&'X'));
247        assert_eq!(tree.get(&2), Some(&'Y'));
248        assert_eq!(tree.get(&0), Some(&'Z'));
249    }
250
251    #[test]
252    fn index() {
253        let mut tree = BinaryTree::new();
254        tree.insert(0, 'A');
255        assert_eq!(tree[&0], 'A');
256    }
257
258    #[test]
259    #[should_panic]
260    fn index_nonexistent() {
261        let mut tree = BinaryTree::new();
262        tree.insert(0, 'A');
263        let _ = tree[&1];
264    }
265}
collect_me/tree/binary_tree.rs

collect_me/tree/
binary_tree.rs
