Skip to main content

Project2/
avl_tree.rs

1// Our code base is adapted from: https://play.rust-lang.org/?gist=d65d605a48d38648737ad2ae38f46434&version=stable
2
3use slab::Slab;
4use std::fmt;
5use std::ops::{Index, IndexMut};
6use std::cmp;
7extern crate slab;
8
9impl<T: fmt::Debug + Copy + fmt::Debug> fmt::Debug for AVLTree<T> {
10    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
11
12        // recursivley print the tree in an ordered fashion
13        fn write_recursive<T: fmt::Debug + Copy>(avltree: &AVLTree<T>, node: Pointer, f: &mut fmt::Formatter){
14            if node.is_null(){
15                write!(f, "").unwrap();
16            }
17            else{
18                write_recursive(avltree, avltree[node].left, f);
19                let left = avltree[node].left;
20                let right = avltree[node].right;
21                let parent = avltree[node].parent;
22
23                write!(f, "(value = {:?},", avltree[node].value).unwrap();
24                
25                if left.is_null(){
26                    write!(f, "left = NULL, ").unwrap();
27                }
28                else{
29                    write!(f, "left = {:?}, ", avltree[left].value).unwrap();
30                }
31
32                if right.is_null(){
33                    write!(f, "right = NULL, ").unwrap();
34                }
35                else{
36                    write!(f, "right = {:?}, ", avltree[right].value).unwrap();
37                }
38
39                if parent.is_null(){
40                    write!(f, "parent = NULL").unwrap();
41                }
42                else{
43                    write!(f, "parent = {:?}", avltree[parent].value).unwrap();
44                }
45
46                write!(f, "), \n").unwrap();
47
48                write_recursive(avltree, avltree[node].right, f);
49
50            }
51        }
52
53        write!(f, "In order traversal:(\n")?;
54        write_recursive(&self, self.root, f);
55        write!(f, ")")?;
56        
57        Ok(())
58    }
59}
60
61#[derive(Eq, PartialEq, Copy, Clone)]
62pub struct Pointer(usize);
63impl Pointer {
64    #[inline]
65    fn null() -> Pointer {
66        Pointer(!0)
67    }
68    
69    #[inline]
70    fn is_null(&self) -> bool {
71        *self == Pointer::null()
72    }
73}
74
75// Just for convenience, so that we can type `self[i]` instead of `self.slab[i]`.
76impl<T> Index<Pointer> for AVLTree<T> {
77    type Output = Node<T>;
78    
79    fn index(&self, index: Pointer) -> &Node<T> {
80        &self.slab[index.0]
81    }
82}
83// Just for convenience, so that we can type `self[i]` instead of `self.slab[i]`.
84impl<T> IndexMut<Pointer> for AVLTree<T> {
85    fn index_mut(&mut self, index: Pointer) -> &mut Node<T> {
86        &mut self.slab[index.0]
87    }
88}
89
90pub struct Node<T> {
91    pub value: T,
92    right: Pointer,
93    left: Pointer,
94    parent: Pointer,
95}
96
97pub struct AVLTree<T> {
98    slab: Slab<Node<T>>,
99    pub root: Pointer,
100}
101
102impl<T: PartialOrd + Copy + fmt::Debug> AVLTree<T> {
103    // Returns a new doubly linked list.
104    pub fn new() -> Self {
105        AVLTree {
106            slab: Slab::new(),
107            root: Pointer::null(),
108        }
109    }
110
111    // Returns true if tree is empty, false otherwise
112    pub fn is_empty(&self) -> bool{
113        return self.root.is_null();
114    }
115
116    // Returns height of tree
117    pub fn get_height(&self) -> u32{
118        return self.get_height_from_node(self.root);
119    }
120
121    // Returns balance factor used to determine balance of AVL tree. Neg number = left heavy, pos number = right heavy
122    pub fn get_balance_factor(&self, node: Pointer) -> i32{
123        return self.get_height_from_node(self[node].right) as i32 - self.get_height_from_node(self[node].left) as i32;
124    }
125
126    // Returns height below node passed as argument
127    pub fn get_height_from_node(&self, node: Pointer) -> u32{
128        if node.is_null(){
129            return 0;
130        }
131        else{
132            let left = self.get_height_from_node(self[node].left);
133            let right = self.get_height_from_node(self[node].right);
134            return cmp::max(left, right) + 1;
135        }
136    }
137
138    // Insert node with value val into tree
139    pub fn insert(&mut self, val: T){
140        if self.root.is_null(){
141            self.root = Pointer(self.slab.insert(Node {
142                value: val,
143                right: Pointer::null(),
144                left: Pointer::null(),
145                parent: Pointer::null(),
146            }));
147        }
148        else{
149            self.insert_below_node(val, self.root);
150        }
151    }
152
153    pub fn insert_below_node(&mut self, val: T, node: Pointer){
154        let nodeValue = self[node].value;
155        let left = self[node].left;
156        let right = self[node].right;
157
158        if val == nodeValue{
159            println!("Duplicate node values");
160            return;
161        }
162        else if val > nodeValue{
163            if right.is_null(){
164                self[node].right = Pointer(self.slab.insert(Node {
165                    value: val,
166                    right: Pointer::null(),
167                    left: Pointer::null(),
168                    parent: node,
169                }));
170                // Fix tree
171                self.rebalance(node);
172            }
173            else{
174                self.insert_below_node(val, right);
175            }
176        }
177        else if val < nodeValue{
178            if left.is_null(){
179                self[node].left = Pointer(self.slab.insert(Node {
180                    value: val,
181                    right: Pointer::null(),
182                    left: Pointer::null(),
183                    parent: node,
184                }));
185                // Fix tree
186                self.rebalance(node);
187            }
188            else{
189                self.insert_below_node(val, left);
190            }
191        }
192    }
193
194    pub fn left_rotate(&mut self, current: Pointer){
195        let right = self[current].right;
196
197        if right.is_null(){
198            println!("NULL");
199            return;
200        }
201
202        let right_left = self[right].left;
203        let parent = self[current].parent;
204
205        // set W's right child to be B
206        self[current].right = right_left;
207
208        if !right_left.is_null(){
209            self[right_left].parent = current;
210        }
211
212        // setting W's parent to be V
213        self[current].parent = right;
214        self[right].left = current;
215
216        // Set V's parent to be W's old parent
217        self[right].parent = parent;
218
219        if parent.is_null(){
220            self.root = right;
221        }
222        else{
223            let parent_right = self[parent].right;
224            if !parent_right.is_null(){
225                if self[parent_right].value == self[current].value{ // set V to parent right
226                    self[parent].right = right;
227                }
228                else{ // set V to parent left
229                    self[parent].left = right;
230                }
231            }
232            else{ // set V to parent left
233                self[parent].left = right;
234            }
235        }
236    }
237    pub fn right_rotate(&mut self, current: Pointer){
238        let left = self[current].left;
239
240        if left.is_null(){
241            return;
242        }
243
244        let left_right = self[left].right;
245        let parent = self[current].parent;
246
247        // set V's left child to be B
248        self[current].left = left_right;
249
250        if !left_right.is_null(){
251            self[left_right].parent = current;
252        }
253
254        // setting V's parent to be W
255        self[current].parent = left;
256        self[left].right = current;
257
258        // Set W's parent to be V's old parent
259        self[left].parent = parent;
260
261        if parent.is_null(){
262            self.root = left;
263        }
264        else{
265            let parent_left = self[parent].left;
266            if !parent_left.is_null(){
267                if self[parent_left].value == self[current].value{ // set W to parent left
268                    self[parent].left = left;
269                }
270                else{ // set W to parent right
271                    self[parent].right = left;
272                }
273            }
274            else{ // set W to parent right
275                self[parent].right = left;
276            }
277        }
278    }
279
280    // Rebalance to ensure AVL tree properties are maintained
281    pub fn rebalance(&mut self, mut node: Pointer){
282        while !node.is_null(){
283            let bal = self.get_balance_factor(node);
284            if bal < -1{                
285                // Left heavy so rotate right
286                let y = self[node].left;
287                let bal_y = self.get_balance_factor(y);
288                if bal_y > 0 {
289                    // Need left-right rotate
290                    self.left_rotate(y);
291                }
292                self.right_rotate(node);
293            }
294            else if bal > 1{
295                // Right heavy so rotate left
296                let y = self[node].right;
297                let bal_y = self.get_balance_factor(y);
298                if bal_y < 0 {
299                    // Need right-left rortate
300                    self.right_rotate(y);
301                }
302                self.left_rotate(node);
303            }
304            node = self[node].parent;
305        }
306    }
307
308    pub fn get_node(&self, val: T) -> Pointer{
309        let node = self.get_node_from_node(self.root, val);
310
311        if node.is_null(){
312            println!("Node does not exist!");
313            return Pointer::null();
314        }
315        return node;
316    }
317
318    pub fn get_node_from_node(&self, node: Pointer, val:T) -> Pointer{
319        if node.is_null(){
320            return Pointer::null();
321        }
322        else{
323            if self[node].value == val{
324                return node;
325            }
326            else if val > self[node].value{
327                return self.get_node_from_node(self[node].right, val);
328            }
329            else{
330                return self.get_node_from_node(self[node].left, val);
331            }
332        }
333    }
334
335    // Delete node with value val
336    pub fn delete(&mut self, val: T) /*-> T*/{
337        let remove = self.get_node(val);
338        if remove.is_null(){
339            return;
340        }
341        let parent = self[remove].parent;
342        // Three cases no children, 1 children, 2 children
343        if self[remove].left.is_null() && self[remove].right.is_null(){
344            // No children just delete node
345            if parent.is_null(){
346                self.root = Pointer::null();
347            }
348            else if self[self[remove].parent].left == remove{
349                self[parent].left = Pointer::null();
350            }
351            else{
352                self[parent].right = Pointer::null();
353            }
354        }
355        else if !self[remove].left.is_null() && !self[remove].right.is_null(){
356            // Two childre need to find replacement node
357            let replace = self.min_of_right(remove);
358            if parent.is_null(){
359                let lefttree = self[remove].left;
360                self.root = replace;
361                self[replace].parent = Pointer::null();
362                self[replace].left = self[remove].left;
363                self[lefttree].parent = replace;
364                if self[remove].right == replace{
365                    self[replace].right = Pointer::null();
366                }
367                else{
368                    let righttree = self[remove].right;
369                    self[replace].right = self[remove].right;
370                    self[righttree].parent = replace;
371                }
372            }
373            else if self[parent].left == remove{
374                let lefttree = self[remove].left;
375                self[parent].left = replace;
376                self[replace].parent = parent;
377                self[replace].left = self[remove].left;
378                self[lefttree].parent = replace;
379                if self[remove].right == replace{
380                    self[replace].right = Pointer::null();
381                }
382                else{
383                    let righttree = self[remove].right;
384                    self[replace].right = self[remove].right;
385                    self[righttree].parent = replace;
386                }
387            }
388            else{
389                let lefttree = self[remove].left;
390                self[parent].right = replace;
391                self[replace].parent = parent;
392                self[replace].left = self[remove].left;
393                self[lefttree].parent = replace;
394                if self[remove].right == replace{
395                    self[replace].right = Pointer::null();
396                }
397                else{
398                    let righttree = self[remove].right;
399                    self[replace].right = self[remove].right;
400                    self[righttree].parent = replace;
401                }
402            }
403        }
404        else{
405            // One child, replace remove with child
406            if parent.is_null(){
407                if self[remove].left.is_null(){
408                    let right = self[remove].right;
409                    self.root = right;
410                    self[right].parent = Pointer::null();
411                }
412                else{
413                    let left = self[remove].left;
414                    self.root = self[remove].left;      
415                    self[left].parent = Pointer::null();
416                }
417            }
418            else if !self[remove].left.is_null(){
419                if self[self[remove].parent].left == remove{
420                    let left = self[remove].left;
421                    self[parent].left = left;
422                    self[left].parent = parent;
423                }
424                else{
425                    let left = self[remove].left;
426                    self[parent].right = left;
427                    self[left].parent = parent;
428                }
429            }
430            else{
431                if self[self[remove].parent].left == remove{
432                    let right = self[remove].right;
433                    self[parent].left = right;
434                    self[right].parent = parent;
435                }
436                else{
437                    let right = self[remove].right;
438                    self[parent].right = right;
439                    self[right].parent = parent;
440                }
441            }
442        }        
443        self.rebalance(parent);
444    }
445
446    // Return number of leaf nodes is tree
447    pub fn count_leaf_nodes(&self) -> u32{
448        return self.count_leaf_nodes_from_node(self.root);
449    }
450    pub fn count_leaf_nodes_from_node(&self, node: Pointer) -> u32{
451        if node.is_null(){
452            return 0;
453        }
454        else{
455            let left = self.count_leaf_nodes_from_node(self[node].left);
456            let right = self.count_leaf_nodes_from_node(self[node].right);
457            if left == right && left == 0{
458                return 1;
459            }
460            return left + right;
461        }
462    }
463
464    // Print tree from left to right
465    pub fn print_in_order_traversal(&self){
466        println!("{:?}", self);
467    }
468
469    // Find smallest value in right tree of head, used for delete
470    pub fn min_of_right(&self, head: Pointer)-> Pointer{
471        let mut current = self[head].right;
472        if current.is_null(){
473            return current;
474        }
475
476        while !self[current].left.is_null(){
477            current = self[current].left;
478        }
479        return current;
480    }
481
482}