tree_collections/
avlTree.rs

1//! AVL tree
2//!
3//! You can generate an AVL tree, and insert or delete nodes.
4
5use std::cell::RefCell;
6use std::fmt::{Debug, Display};
7use std::rc::Rc;
8
9pub use crate::commonTrait::{CommonTreeNodeTrait, CommonTreeTrait};
10
11type AVLTreeNode<T> = Rc<RefCell<TreeNode<T>>>;
12type OptionAVLTreeNode<T> = Option<AVLTreeNode<T>>;
13
14/// Node struct for AVLTree
15#[derive(Clone, Debug, PartialEq)]
16pub struct TreeNode<T: Ord + Copy + Debug + Display> {
17    pub value: T,
18    left: OptionAVLTreeNode<T>,
19    right: OptionAVLTreeNode<T>,
20    height: usize,
21}
22
23// extend from common tree trait
24impl<T: Ord + Copy + Debug + Display> CommonTreeTrait<T, TreeNode<T>> for AVLTree<T> {
25    fn get_root(&self) -> OptionAVLTreeNode<T> {
26        return self.root.clone();
27    }
28}
29
30// extend from common tree node trait
31impl<T: Ord + Copy + Debug + Display> CommonTreeNodeTrait<T> for TreeNode<T> {
32    fn get_left(&self) -> OptionAVLTreeNode<T> {
33        return self.left.clone();
34    }
35
36    fn get_right(&self) -> OptionAVLTreeNode<T> {
37        return self.right.clone();
38    }
39
40    fn get_value(&self) -> T {
41        return self.value;
42    }
43
44    fn get_value_to_print(&self) -> String {
45        return self.value.to_string();
46    }
47}
48
49/// Implementations of AVLTreeNode
50impl<T: Ord + Copy + Debug + Display> TreeNode<T> {
51    /// Create a new node of type OptionAVLTreeNode , which will be called by [AVLTree](struct.AVLTree.html)
52    fn new(value: T) -> OptionAVLTreeNode<T> {
53        Some(Rc::new(RefCell::new(Self {
54            value,
55            left: None,
56            right: None,
57            height: 1, // default height of a new node is 1，which is a leave
58        })))
59    }
60}
61
62#[derive(Clone, Debug, PartialEq)]
63pub struct AVLTree<T: Ord + Copy + Debug + Display> {
64    root: OptionAVLTreeNode<T>,
65}
66
67/// Implementations of AVLTree
68impl<T: Ord + Copy + Debug + Display> AVLTree<T> {
69    /// Creates a new AVL tree
70    ///
71    /// # Example
72    ///
73    /// ```
74    /// use tree_collections::avlTree::AVLTree;
75    /// let mut avl_tree: AVLTree<u32> = AVLTree::new();
76    /// ```
77    pub fn new() -> Self {
78        Self { root: None }
79    }
80
81    pub fn preorder_traverse(&self, node: AVLTreeNode<T>, container: &mut Vec<T>) {
82        container.push(node.borrow().value);
83        let left = node.borrow().left.clone();
84        if left.is_some() {
85            self.preorder_traverse(left.unwrap(), container);
86        }
87        let right = node.borrow().right.clone();
88        if right.is_some() {
89            self.preorder_traverse(right.unwrap(), container);
90        }
91    }
92
93    pub fn in_order_traverse(&self, node: AVLTreeNode<T>, container: &mut Vec<T>) {
94        let left = node.borrow().left.clone();
95        if left.is_some() {
96            self.in_order_traverse(left.unwrap(), container);
97        }
98        container.push(node.borrow().value);
99        let right = node.borrow().right.clone();
100        if right.is_some() {
101            self.in_order_traverse(right.unwrap(), container);
102        }
103    }
104
105    /// Judge if the AVL tree is empty
106    ///
107    /// # Example
108    ///
109    /// ```
110    /// use tree_collections::avlTree::AVLTree;
111    /// let mut avl_tree = AVLTree::new();
112    /// println!("{}", avl_tree.is_tree_empty());  // true
113    /// avl_tree.insert(1);
114    /// println!("{}", avl_tree.is_tree_empty());  // false
115    /// ```
116    pub fn is_tree_empty(&self) -> bool {
117        self.root.clone().map(|_| false).unwrap_or(true)
118    }
119
120    pub fn insert(&mut self, insert_value: T) {
121        let root = self.root.take();
122        // TreeNode is type OptionAVLTreeNode, so the code is simplified.
123        match root {
124            None => self.root = TreeNode::new(insert_value),
125            Some(n) => self.root = self.node_insert(Some(n), insert_value),
126        }
127    }
128
129    /// Inserts a node, return a new root, which will be called by
130    /// [AVLTree.insert](struct.AVLTree.html#method.insert)
131    fn node_insert(&mut self, node: OptionAVLTreeNode<T>, insert_value: T) -> OptionAVLTreeNode<T> {
132        let ret_node = match node {
133            Some(n) => {
134                let node_value = n.borrow().value;
135                if insert_value < node_value {
136                    let left = n.borrow().left.clone();
137                    n.borrow_mut().left = self.node_insert(left, insert_value);
138                } else if insert_value > node_value {
139                    let right = n.borrow().right.clone();
140                    n.borrow_mut().right = self.node_insert(right, insert_value);
141                } else {
142                    n.borrow_mut().value = insert_value; // equal, update value
143                }
144                n
145            }
146            None => TreeNode::new(insert_value).unwrap(),
147        };
148
149        // update height
150        ret_node.borrow_mut().height = self
151            .get_left_height(&ret_node)
152            .max(self.get_right_height(&ret_node))
153            + 1;
154
155        // update balance factor
156        let balance_factor = self.get_balance_factor(&ret_node);
157
158        // maintain
159        // case LL: right rotate
160        if balance_factor > 1.0
161            && self.get_balance_factor(&ret_node.borrow().left.clone().unwrap()) >= 0.0
162        {
163            return Some(self.right_rotate(ret_node));
164        }
165
166        // case RR: left rotate
167        if balance_factor < -1.0
168            && self.get_balance_factor(&ret_node.borrow().right.clone().unwrap()) <= 0.0
169        {
170            return Some(self.left_rotate(ret_node));
171        }
172
173        // case LR: left rotate + right rotate
174        if balance_factor > 1.0
175            && self.get_balance_factor(&ret_node.borrow().left.clone().unwrap()) < 0.0
176        {
177            // ret_node.borrow_mut().left = Some(self.left_rotate(ret_node.borrow_mut().left.clone().unwrap())); // 发生移动
178            // return Some(self.right_rotate(ret_node))
179
180            let left = ret_node.borrow().left.clone().take().unwrap();
181            ret_node.borrow_mut().left = Some(self.left_rotate(left));
182            return Some(self.right_rotate(ret_node));
183        }
184
185        // case RL: right rotate + left rotate
186        if balance_factor < -1.0
187            && self.get_balance_factor(&ret_node.borrow().right.clone().unwrap()) > 0.0
188        {
189            // ret_node.borrow_mut().right = Some(self.right_rotate(ret_node.borrow_mut().right.clone().unwrap())); // 发生移动
190            // return Some(self.left_rotate(ret_node))
191
192            let right = ret_node.borrow().right.clone().take().unwrap();
193            ret_node.borrow_mut().right = Some(self.right_rotate(right));
194            return Some(self.left_rotate(ret_node));
195        }
196        Some(ret_node)
197    }
198
199    /// Delete a value from AVL tree
200    ///
201    /// # Example
202    ///
203    /// ```
204    /// use tree_collections::avlTree::AVLTree;
205    /// let mut avl_tree = AVLTree::new();
206    /// avl_tree.insert(1);
207    /// avl_tree.delete(1);
208    /// ```
209    pub fn delete(&mut self, delete_value: T) {
210        let root = self.root.take();
211        match root {
212            None => return,
213            Some(n) => self.root = self.node_delete(Some(n), delete_value),
214        }
215    }
216    /// Deletes a node, return a new root, which will be called by
217    /// [AVLTree.delete](struct.AVLTree.html#method.delete)
218    // delete node, return new root
219    fn node_delete(&mut self, node: OptionAVLTreeNode<T>, delete_value: T) -> OptionAVLTreeNode<T> {
220        let ret_node = match node {
221            None => node,
222            Some(mut n) => {
223                let node_value = n.borrow().value;
224                if delete_value < node_value {
225                    // look left
226                    let left = n.borrow().left.clone();
227                    n.borrow_mut().left = self.node_delete(left, delete_value);
228                    Some(n)
229                } else if delete_value > node_value {
230                    // look right
231                    let right = n.borrow().right.clone();
232                    n.borrow_mut().right = self.node_delete(right, delete_value);
233                    Some(n)
234                } else {
235                    // found the node which should be deleted
236                    let left = n.borrow().left.clone();
237                    let right = n.borrow().right.clone();
238                    let ret = match (left.clone(), right.clone()) {
239                        (None, Some(r)) => Some(r), // The left subtree of the node to be deleted is empty, r is new root
240                        (Some(l), None) => Some(l), // The right subtree of the node to be deleted is empty, l is new root
241
242                        (None, None) => None,
243
244                        // The left and right subtrees of the node to be deleted(node n) are not empty.
245                        // Find the smallest node A that is larger than the node n.
246                        (Some(_), Some(right)) => {
247                            let min_value = right.borrow().get_min_value_in_children(); // Find the value of node A which is the minimum value of the right subtree
248                            n.borrow_mut().value = min_value; // Change the value of node n to the value of node A.
249                            let right = n.borrow().right.clone().take();
250                            n.borrow_mut().right = self.node_delete(right, min_value); // Delete the node A in the right subtree.
251                            Some(n) // return new root
252                        }
253                    };
254                    ret // 返回option
255                }
256            }
257        };
258
259        // update and maintain
260        match ret_node {
261            None => ret_node,
262            Some(n) => {
263                // update height
264                n.borrow_mut().height = self
265                    .get_left_height(&n) // 借用了发生移动的
266                    .max(self.get_right_height(&n))
267                    + 1; // 把option类型的ret_node都改成了n
268
269                // update balance factor
270                let balance_factor = self.get_balance_factor(&n);
271
272                // maintain
273                // case LL: right rotate
274                if balance_factor > 1.0
275                    && self.get_balance_factor(&n.borrow().left.clone().unwrap()) >= 0.0
276                {
277                    return Some(self.right_rotate(n));
278                }
279
280                // case RR: left rotate
281                if balance_factor < -1.0
282                    && self.get_balance_factor(&n.borrow().right.clone().unwrap()) <= 0.0
283                {
284                    return Some(self.left_rotate(n));
285                }
286
287                // case LR: left rotate + right rotate
288                if balance_factor > 1.0
289                    && self.get_balance_factor(&n.borrow().left.clone().unwrap()) < 0.0
290                {
291                    let left = n.borrow().left.clone().take().unwrap();
292                    n.borrow_mut().left = Some(self.left_rotate(left));
293                    return Some(self.right_rotate(n));
294                }
295
296                // case RL: right rotate + left rotate
297                if balance_factor < -1.0
298                    && self.get_balance_factor(&n.borrow().right.clone().unwrap()) > 0.0
299                {
300                    let right = n.borrow().right.clone().take().unwrap();
301                    n.borrow_mut().right = Some(self.right_rotate(right));
302                    return Some(self.left_rotate(n));
303                }
304                Some(n)
305            }
306        }
307    }
308
309    fn get_height(&self, node: OptionAVLTreeNode<T>) -> usize {
310        // default height of an empty tree is 0
311        node.map_or(0, |n| n.borrow().height)
312    }
313    fn get_left_height(&self, n: &AVLTreeNode<T>) -> usize {
314        self.get_height(n.borrow().left.clone())
315    }
316
317    fn get_right_height(&self, n: &AVLTreeNode<T>) -> usize {
318        self.get_height(n.borrow().right.clone())
319    }
320
321    fn get_balance_factor(&self, n: &AVLTreeNode<T>) -> f64 {
322        self.get_left_height(n) as f64 - self.get_right_height(n) as f64
323    }
324
325    //Determine whether the tree is balanced
326    fn is_balanced(&self, node: OptionAVLTreeNode<T>) -> bool {
327        match node {
328            Some(node) => {
329                if self.get_balance_factor(&node) <= 1.0 {
330                    self.is_balanced(node.borrow().left.clone())
331                        && self.is_balanced(node.borrow().right.clone())
332                } else {
333                    false
334                }
335            }
336            None => true,
337        }
338    }
339
340    //                 y                                     x
341    //               /    \                                 /   \
342    //              x     T4      right rotate (y)         z     y
343    //           /   \            ---------------->       / \   / \
344    //          z     T3             return x            T1 T2 T3 T4
345    //         /   \
346    //        T1   T2
347    fn right_rotate(&self, y: AVLTreeNode<T>) -> AVLTreeNode<T> {
348        let x = y.borrow().left.clone().unwrap();
349        let t_3 = x.borrow().right.clone().take();
350
351        // right rotate
352        x.borrow_mut().right = Some(y.clone());
353        y.borrow_mut().left = t_3;
354
355        // update height of x and y
356        y.borrow_mut().height = self.get_left_height(&y).max(self.get_right_height(&y)) + 1;
357        x.borrow_mut().height = self.get_left_height(&x).max(self.get_right_height(&x)) + 1;
358
359        return x;
360    }
361
362    //                 y                                        x
363    //               /    \                                   /   \
364    //              T1     x        left rotate (y)          y     z
365    //                   /   \      ---------------->       / \   / \
366    //                 T2     z        return x            T1 T2 T3 T4
367    //                       /  \
368    //                      T3   T4
369    fn left_rotate(&self, y: AVLTreeNode<T>) -> AVLTreeNode<T> {
370        let x = y.borrow().right.clone().unwrap();
371        let t_2 = x.borrow().left.clone().take();
372
373        // left rotate
374        x.borrow_mut().left = Some(y.clone());
375        y.borrow_mut().right = t_2;
376
377        // update height of x and y
378        y.borrow_mut().height = self.get_left_height(&y).max(self.get_right_height(&y)) + 1;
379        x.borrow_mut().height = self.get_left_height(&x).max(self.get_right_height(&x)) + 1;
380
381        return x;
382    }
383}
384
385impl<T: Ord + Copy + Debug + Display> Drop for AVLTree<T> {
386    fn drop(&mut self) {
387        match self.root.take() {
388            Some(node) => node.borrow_mut().clear(),
389            None => return,
390        }
391    }
392}
393
394impl<T: Ord + Copy + Debug + Display> Drop for TreeNode<T> {
395    fn drop(&mut self) {
396        self.clear();
397    }
398}
399
400impl<T: Ord + Copy + Debug + Display> TreeNode<T> {
401    fn clear(&mut self) {
402        match self.left.take() {
403            None => {}
404            Some(node) => {
405                node.borrow_mut().clear();
406            }
407        }
408        self.left = None;
409        match self.right.take() {
410            None => {}
411            Some(node) => {
412                node.borrow_mut().clear();
413            }
414        }
415        self.right = None;
416    }
417}
418
419#[cfg(test)]
420mod test {
421    use super::*;
422
423    #[test]
424    fn tree_traversal() {
425        // Test the three different tree traversal functions.
426        let mut tree = AVLTree::new();
427        tree.insert(0);
428        vec![16, 16, 8, 24, 20, 22].iter().for_each(|v| {
429            tree.insert(*v);
430        });
431        let root = tree.root.clone().unwrap();
432        let mut pre_container = vec![];
433        let mut in_container = vec![];
434        tree.preorder_traverse(root.clone(), &mut pre_container);
435        tree.in_order_traverse(root.clone(), &mut in_container);
436        let is_balanced = tree.is_balanced(tree.root.clone());
437        // println!("check {:#?}", in_container);
438        assert_eq!(pre_container, vec![20, 8, 0, 16, 24, 22]);
439        assert_eq!(in_container, vec![0, 8, 16, 20, 22, 24]);
440        assert_eq!(is_balanced, true);
441    }
442
443    #[test]
444    fn test_insert() {
445        let mut avl_tree = AVLTree::new();
446        avl_tree.insert(1);
447        avl_tree.insert(2);
448        avl_tree.insert(3);
449        avl_tree.insert(4);
450        avl_tree.insert(5);
451
452        let result = avl_tree.is_balanced(avl_tree.root.clone());
453        assert_eq!(result, true);
454    }
455
456    #[test]
457    fn test_delete() {
458        // Test the three different tree traversal functions.
459        let mut tree = AVLTree::new();
460        tree.insert(0);
461        vec![16, 8, 24, 20, 22].iter().for_each(|v| {
462            tree.insert(*v);
463        });
464
465        let root = tree.root.clone().unwrap();
466        tree.delete(16);
467        let mut container = vec![];
468        tree.preorder_traverse(root.clone(), &mut container);
469        let result = tree.is_balanced(tree.root.clone());
470        assert_eq!(result, true);
471
472        assert_eq!(container, vec![20, 8, 0, 24, 22]);
473    }
474}
tree_collections/avlTree.rs

tree_collections/
avlTree.rs
