Parcourir la source

feat(node): add basic implementation of AVL nodes

Danilo Gómez il y a 1 an
Parent
commit
21edc3bcba
1 fichiers modifiés avec 97 ajouts et 0 suppressions
  1. 97 0
      src/main.rs

+ 97 - 0
src/main.rs

@@ -1 +1,98 @@
+use std::{
+    cmp::{max, Ordering},
+    mem::swap,
+};
+
+struct Node<T: Ord> {
+    value: T,
+    height: i32,
+    left: Option<Box<Node<T>>>,
+    right: Option<Box<Node<T>>>,
+}
+
+impl<T: Ord> Node<T> {
+    fn new(value: T) -> Self {
+        Node {
+            value,
+            height: 1,
+            left: None,
+            right: None,
+        }
+    }
+
+    fn balance_factor(&self) -> i32 {
+        let lheight = self.left.as_ref().map_or(0, |n| n.height);
+        let rheight = self.right.as_ref().map_or(0, |n| n.height);
+        lheight - rheight
+    }
+
+    fn update_height(&mut self) {
+        let lheight = self.left.as_ref().map_or(0, |n| n.height);
+        let rheight = self.right.as_ref().map_or(0, |n| n.height);
+        self.height = max(lheight, rheight) + 1;
+    }
+
+    fn rotate_right(&mut self) {
+        let mut x = self.left.take().unwrap();
+        self.left = x.right.take();
+        self.update_height();
+        swap(self, &mut x);
+        self.right = Some(x);
+        self.update_height();
+    }
+
+    fn rotate_left(&mut self) {
+        let mut x = self.right.take().unwrap();
+        self.right = x.left.take();
+        self.update_height();
+        swap(self, &mut x);
+        self.left = Some(x);
+        self.update_height();
+    }
+
+    fn balance(&mut self) {
+        if self.balance_factor() > 1 {
+            if self.left.as_ref().unwrap().balance_factor() < 0 {
+                let mut left = self.left.take().unwrap();
+                left.rotate_left();
+                self.left = Some(left);
+            }
+            self.rotate_right();
+        } else if self.balance_factor() < -1 {
+            if self.right.as_ref().unwrap().balance_factor() > 0 {
+                let mut right = self.right.take().unwrap();
+                right.rotate_right();
+                self.right = Some(right);
+            }
+            self.rotate_left();
+        }
+    }
+
+    fn insert(&mut self, value: T) -> bool {
+        let ret_val = match value.cmp(&self.value) {
+            Ordering::Less => {
+                if let Some(ref mut left) = self.left {
+                    left.insert(value)
+                } else {
+                    self.left = Some(Box::new(Node::new(value)));
+                    true
+                }
+            }
+            Ordering::Greater => {
+                if let Some(ref mut right) = self.right {
+                    right.insert(value)
+                } else {
+                    self.right = Some(Box::new(Node::new(value)));
+                    true
+                }
+            }
+            Ordering::Equal => false,
+        };
+
+        self.update_height();
+        self.balance();
+        ret_val
+    }
+}
+
 fn main() {}