range_filters/
binary_search_tree.rs

1use crate::Key;
2use crate::infix_store::InfixStore;
3use std::fmt;
4use std::sync::{Arc, RwLock};
5
6// #[derive(Debug, Default)]
7// pub struct InfixStore;
8
9// TODO: add cached count
10#[derive(Debug, Default)]
11pub struct BinarySearchTreeGroup {
12    pub root: Option<Box<TreeNode>>,
13}
14
15#[derive(Clone, Debug)]
16pub struct TreeNode {
17    pub key: Key,
18    pub left: Option<Box<TreeNode>>,
19    pub right: Option<Box<TreeNode>>,
20    pub infix_store: Option<Arc<RwLock<InfixStore>>>,
21}
22
23impl BinarySearchTreeGroup {
24    pub fn new() -> Self {
25        Self { root: None }
26    }
27
28    pub fn new_with_keys(keys: &[Key]) -> Self {
29        if keys.is_empty() {
30            return Self { root: None };
31        }
32
33        let mut sorted_keys = keys.to_vec();
34        sorted_keys.sort();
35
36        let root = Self::top_down_bst_insertion(&sorted_keys, 0, sorted_keys.len() as isize - 1);
37        Self { root }
38    }
39
40    fn top_down_bst_insertion(keys: &[Key], start: isize, end: isize) -> Option<Box<TreeNode>> {
41        if start > end {
42            return None;
43        }
44
45        let mid = ((start + end) / 2) as usize;
46        let root = Box::new(TreeNode {
47            key: keys[mid],
48            left: Self::top_down_bst_insertion(keys, start, mid as isize - 1),
49            right: Self::top_down_bst_insertion(keys, mid as isize + 1, end),
50            infix_store: None,
51        });
52        Some(root)
53    }
54
55    // TODO: use cached length
56    pub fn len(&self) -> usize {
57        Self::len_recursive(&self.root)
58    }
59
60    fn len_recursive(node: &Option<Box<TreeNode>>) -> usize {
61        match node {
62            None => 0,
63            Some(n) => 1 + Self::len_recursive(&n.left) + Self::len_recursive(&n.right),
64        }
65    }
66
67    pub fn insert(&mut self, key: Key) {
68        Self::insert_recursive(&mut self.root, key);
69    }
70
71    fn insert_recursive(node: &mut Option<Box<TreeNode>>, key: Key) {
72        match node {
73            None => {
74                *node = Some(Box::new(TreeNode {
75                    key,
76                    left: None,
77                    right: None,
78                    infix_store: None,
79                }));
80            }
81            Some(n) => {
82                if key < n.key {
83                    Self::insert_recursive(&mut n.left, key);
84                } else {
85                    Self::insert_recursive(&mut n.right, key);
86                }
87            }
88        }
89    }
90
91    pub fn contains(&self, key: Key) -> bool {
92        Self::contains_recursive(&self.root, key)
93    }
94
95    fn contains_recursive(node: &Option<Box<TreeNode>>, key: Key) -> bool {
96        match node {
97            None => false,
98            Some(n) => {
99                if key == n.key {
100                    true
101                } else if key < n.key {
102                    Self::contains_recursive(&n.left, key)
103                } else {
104                    Self::contains_recursive(&n.right, key)
105                }
106            }
107        }
108    }
109
110    fn find_node_mut(node: &mut Option<Box<TreeNode>>, key: Key) -> Option<&mut TreeNode> {
111        match node {
112            None => None,
113            Some(n) => {
114                if key == n.key {
115                    Some(n.as_mut())
116                } else if key < n.key {
117                    Self::find_node_mut(&mut n.left, key)
118                } else {
119                    Self::find_node_mut(&mut n.right, key)
120                }
121            }
122        }
123    }
124
125    pub fn set_infix_store(&mut self, key: Key, infix_store: InfixStore) {
126        if let Some(node) = Self::find_node_mut(&mut self.root, key) {
127            node.infix_store = Some(Arc::new(RwLock::new(infix_store)));
128        }
129    }
130
131    pub fn get_infix_store(&self, key: Key) -> Option<Arc<RwLock<InfixStore>>> {
132        Self::get_infix_store_recursive(&self.root, key)
133    }
134
135    fn get_infix_store_recursive(
136        node: &Option<Box<TreeNode>>,
137        key: Key,
138    ) -> Option<Arc<RwLock<InfixStore>>> {
139        match node {
140            None => None,
141            Some(n) => {
142                if key == n.key {
143                    n.infix_store.clone()
144                } else if key < n.key {
145                    Self::get_infix_store_recursive(&n.left, key)
146                } else {
147                    Self::get_infix_store_recursive(&n.right, key)
148                }
149            }
150        }
151    }
152
153    pub fn predecessor_infix_store(&self, key: Key) -> Option<Arc<RwLock<InfixStore>>> {
154        Self::predecessor_store_recursive(&self.root, key, None)
155    }
156
157    pub fn predecessor(&self, key: Key) -> Option<Key> {
158        Self::predecessor_recursive(&self.root, key, None)
159    }
160
161    fn predecessor_recursive(
162        node: &Option<Box<TreeNode>>,
163        key: Key,
164        best: Option<Key>,
165    ) -> Option<Key> {
166        match node {
167            None => best,
168            Some(n) => {
169                if n.key == key {
170                    Some(n.key)
171                } else if key < n.key {
172                    Self::predecessor_recursive(&n.left, key, best)
173                } else {
174                    Self::predecessor_recursive(&n.right, key, Some(n.key))
175                }
176            }
177        }
178    }
179
180    pub fn successor(&self, key: Key) -> Option<Key> {
181        Self::successor_recursive(&self.root, key, None)
182    }
183
184    fn successor_recursive(
185        node: &Option<Box<TreeNode>>,
186        key: Key,
187        best: Option<Key>,
188    ) -> Option<Key> {
189        match node {
190            None => best,
191            Some(n) => {
192                if n.key == key {
193                    Some(n.key)
194                } else if key < n.key {
195                    Self::successor_recursive(&n.left, key, Some(n.key))
196                } else {
197                    Self::successor_recursive(&n.right, key, best)
198                }
199            }
200        }
201    }
202
203    fn predecessor_store_recursive(
204        node: &Option<Box<TreeNode>>,
205        key: Key,
206        best: Option<Arc<RwLock<InfixStore>>>,
207    ) -> Option<Arc<RwLock<InfixStore>>> {
208        match node {
209            None => best,
210            Some(n) => {
211                if n.key == key {
212                    n.infix_store.clone().or(best)
213                } else if key < n.key {
214                    Self::predecessor_store_recursive(&n.left, key, best)
215                } else {
216                    Self::predecessor_store_recursive(&n.right, key, n.infix_store.clone())
217                }
218            }
219        }
220    }
221
222    pub fn successor_infix_store(&self, key: Key) -> Option<Arc<RwLock<InfixStore>>> {
223        Self::successor_store_recursive(&self.root, key, None)
224    }
225
226    fn successor_store_recursive(
227        node: &Option<Box<TreeNode>>,
228        key: Key,
229        best: Option<Arc<RwLock<InfixStore>>>,
230    ) -> Option<Arc<RwLock<InfixStore>>> {
231        match node {
232            None => best,
233            Some(n) => {
234                if n.key == key {
235                    n.infix_store.clone().or(best)
236                } else if key < n.key {
237                    Self::successor_store_recursive(&n.left, key, n.infix_store.clone())
238                } else {
239                    Self::successor_store_recursive(&n.right, key, best)
240                }
241            }
242        }
243    }
244
245    #[allow(dead_code)]
246    fn min_key(node: &Option<Box<TreeNode>>) -> Option<Key> {
247        match node {
248            None => None,
249            Some(n) => {
250                if n.left.is_none() {
251                    Some(n.key)
252                } else {
253                    Self::min_key(&n.left)
254                }
255            }
256        }
257    }
258
259    #[allow(dead_code)]
260    fn max_key(node: &Option<Box<TreeNode>>) -> Option<Key> {
261        match node {
262            None => None,
263            Some(n) => {
264                if n.right.is_none() {
265                    Some(n.key)
266                } else {
267                    Self::max_key(&n.right)
268                }
269            }
270        }
271    }
272
273    #[allow(dead_code)]
274    fn min_node(node: &Option<Box<TreeNode>>) -> Option<&TreeNode> {
275        match node {
276            None => None,
277            Some(n) => {
278                if n.left.is_none() {
279                    Some(n)
280                } else {
281                    Self::min_node(&n.left)
282                }
283            }
284        }
285    }
286
287    #[allow(dead_code)]
288    fn max_node(node: &Option<Box<TreeNode>>) -> Option<&TreeNode> {
289        match node {
290            None => None,
291            Some(n) => {
292                if n.right.is_none() {
293                    Some(n)
294                } else {
295                    Self::max_node(&n.right)
296                }
297            }
298        }
299    }
300
301    pub fn pretty_print(&self) {
302        print!("{}", self);
303    }
304
305    fn format_tree(
306        node: &Option<Box<TreeNode>>,
307        prefix: &str,
308        is_tail: bool,
309        f: &mut fmt::Formatter,
310    ) -> fmt::Result {
311        if let Some(n) = node {
312            writeln!(
313                f,
314                "{}{} {}",
315                prefix,
316                if is_tail { "└──" } else { "├──" },
317                n.key
318            )?;
319
320            let new_prefix = format!("{}{}", prefix, if is_tail { "    " } else { "│   " });
321
322            if n.right.is_some() || n.left.is_some() {
323                if n.right.is_some() {
324                    Self::format_tree(&n.right, &new_prefix, false, f)?;
325                }
326                if n.left.is_some() {
327                    Self::format_tree(&n.left, &new_prefix, true, f)?;
328                }
329            }
330        }
331        Ok(())
332    }
333}
334
335impl fmt::Display for BinarySearchTreeGroup {
336    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
337        writeln!(f, "\n=== Binary Search Tree ===")?;
338        if self.root.is_none() {
339            writeln!(f, "  (empty tree)")?;
340        } else {
341            Self::format_tree(&self.root, "", true, f)?;
342        }
343        writeln!(f, "=========================\n")?;
344        Ok(())
345    }
346}
347
348#[cfg(test)]
349mod tests {
350    use super::*;
351
352    #[test]
353    fn test_tree_construction() {
354        let bst = BinarySearchTreeGroup::new_with_keys(&[1, 2, 3, 20, 30, 4, 5, 6, 7]);
355        assert!(bst.contains(1));
356        assert!(bst.contains(2));
357        assert!(bst.contains(30));
358        assert!(bst.contains(4));
359        assert!(bst.contains(5));
360        assert!(bst.contains(6));
361        assert!(bst.contains(7));
362        assert!(!bst.contains(8));
363        assert!(!bst.contains(9));
364        assert!(!bst.contains(10));
365    }
366
367    #[test]
368    fn test_tree_insertion() {
369        let mut bst = BinarySearchTreeGroup::new();
370        bst.insert(1);
371        bst.insert(2);
372        bst.insert(3);
373        bst.insert(20);
374        bst.insert(30);
375        bst.insert(4);
376        bst.insert(5);
377        bst.insert(6);
378        bst.insert(7);
379        assert!(bst.contains(1));
380        assert!(bst.contains(2));
381        assert!(bst.contains(3));
382        assert!(bst.contains(20));
383        assert!(bst.contains(30));
384        assert!(bst.contains(4));
385        assert!(bst.contains(5));
386        assert!(bst.contains(6));
387        assert!(bst.contains(7));
388        assert!(!bst.contains(8));
389        assert!(!bst.contains(9));
390        assert!(!bst.contains(10));
391    }
392
393    #[test]
394    fn test_predecessor_infix_store() {
395        let mut bst = BinarySearchTreeGroup::new_with_keys(&[10, 20, 30, 40, 50]);
396
397        bst.set_infix_store(10, InfixStore::default());
398        bst.set_infix_store(20, InfixStore::default());
399        bst.set_infix_store(30, InfixStore::default());
400        bst.set_infix_store(40, InfixStore::default());
401        bst.set_infix_store(50, InfixStore::default());
402
403        // exact match returns own store
404        let store_30 = bst.get_infix_store(30).unwrap();
405        let pred_30 = bst.predecessor_infix_store(30).unwrap();
406        assert!(Arc::ptr_eq(&store_30, &pred_30));
407
408        // between keys returns predecessor's store
409        let pred_35 = bst.predecessor_infix_store(35).unwrap();
410        assert!(Arc::ptr_eq(&store_30, &pred_35));
411
412        let store_20 = bst.get_infix_store(20).unwrap();
413        let pred_25 = bst.predecessor_infix_store(25).unwrap();
414        assert!(Arc::ptr_eq(&store_20, &pred_25));
415
416        // before first key
417        assert!(bst.predecessor_infix_store(5).is_none());
418
419        // after last key
420        let store_50 = bst.get_infix_store(50).unwrap();
421        let pred_60 = bst.predecessor_infix_store(60).unwrap();
422        assert!(Arc::ptr_eq(&store_50, &pred_60));
423    }
424}
range_filters/binary_search_tree.rs

range_filters/
binary_search_tree.rs
