Merge pull request #22 from azizkayumov/develop

Clean up splay
azizkayumov · Oct 6, 2023 · aacbf0d · aacbf0d
2 parents 0d9b390 + 0e108e5
commit aacbf0d
Show file tree

Hide file tree

Showing 3 changed files with 95 additions and 44 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -2,9 +2,11 @@
 mod node;
 mod path;
 mod splay;
+use path::update_max;
+
 use crate::{
     node::{Node, Parent},
-    splay::{splay, update_max},
+    splay::splay,
 };
 
 pub struct LinkCutTree {

diff --git a/src/path.rs b/src/path.rs
@@ -41,3 +41,20 @@ impl Path for FindSum {
         self.sum += other.weight;
     }
 }
+
+pub fn update_max(forest: &mut [Node], node_idx: usize) {
+    let mut max_idx = node_idx;
+    if let Some(left_child) = forest[node_idx].left {
+        let left_max = forest[left_child].max_weight_idx;
+        if forest[left_max].weight > forest[max_idx].weight {
+            max_idx = left_max;
+        }
+    }
+    if let Some(right_child) = forest[node_idx].right {
+        let right_max = forest[right_child].max_weight_idx;
+        if forest[right_max].weight > forest[max_idx].weight {
+            max_idx = right_max;
+        }
+    }
+    forest[node_idx].max_weight_idx = max_idx;
+}
diff --git a/src/splay.rs b/src/splay.rs
@@ -1,4 +1,7 @@
-use crate::node::{Node, Parent};
+use crate::{
+    node::{Node, Parent},
+    path::update_max,
+};
 
 /// Rotates the subtree rooted at `node_idx` to the left:
 ///
@@ -84,21 +87,29 @@ fn rotate_right(forest: &mut [Node], node_idx: usize) {
     update_max(forest, left_child);
 }
 
-pub fn update_max(forest: &mut [Node], node_idx: usize) {
-    let mut max_idx = node_idx;
-    if let Some(left_child) = forest[node_idx].left {
-        let left_max = forest[left_child].max_weight_idx;
-        if forest[left_max].weight > forest[max_idx].weight {
-            max_idx = left_max;
-        }
-    }
-    if let Some(right_child) = forest[node_idx].right {
-        let right_max = forest[right_child].max_weight_idx;
-        if forest[right_max].weight > forest[max_idx].weight {
-            max_idx = right_max;
+/// Rotates the parent of `node_idx` to the right or left, depending on the relationship between:
+///
+/// # Panics
+/// Panics if `node_idx` is out of bounds or if `node_idx` does not have a parent.
+///
+/// # Examples:
+///    0                1
+///   /        =>        \
+///  1                    0
+fn rotate(forest: &mut [Node], node_idx: usize) {
+    assert!(node_idx < forest.len(), "rotate: node_idx out of bounds");
+    assert!(
+        matches!(forest[node_idx].parent, Parent::Node(_)),
+        "rotate: node_idx does not have a parent"
+    );
+
+    if let Parent::Node(parent_idx) = forest[node_idx].parent {
+        if forest[parent_idx].left == Some(node_idx) {
+            rotate_right(forest, parent_idx);
+        } else {
+            rotate_left(forest, parent_idx);
         }
     }
-    forest[node_idx].max_weight_idx = max_idx;
 }
 
 /// Splays the subtree rooted at `node_idx`, making it the new root of the tree.
@@ -111,44 +122,29 @@ pub fn update_max(forest: &mut [Node], node_idx: usize) {
 ///    /
 ///   2
 pub fn splay(forest: &mut [Node], node_idx: usize) {
-    assert!(node_idx < forest.len(), "splay: node_idx out of bounds");
-
     while let Parent::Node(parent_idx) = forest[node_idx].parent {
-        if forest[parent_idx].left == Some(node_idx) {
-            if let Parent::Node(grandparent_idx) = forest[parent_idx].parent {
-                if forest[grandparent_idx].left == Some(parent_idx) {
-                    // zig-zig
-                    rotate_right(forest, grandparent_idx);
-                    rotate_right(forest, parent_idx);
-                } else {
-                    // zig-zag
-                    rotate_right(forest, parent_idx);
-                    rotate_left(forest, grandparent_idx);
-                }
-            } else {
-                // zig
-                rotate_right(forest, parent_idx);
-            }
-        } else if let Parent::Node(grandparent_idx) = forest[parent_idx].parent {
-            if forest[grandparent_idx].right == Some(parent_idx) {
-                // zig-zig
-                rotate_left(forest, grandparent_idx);
-                rotate_left(forest, parent_idx);
+        if let Parent::Node(grandparent_idx) = forest[parent_idx].parent {
+            if (forest[grandparent_idx].left == Some(parent_idx))
+                == (forest[parent_idx].left == Some(node_idx))
+            {
+                // zig-zig (same direction):
+                rotate(forest, parent_idx);
+                rotate(forest, node_idx);
             } else {
-                // zig-zag
-                rotate_left(forest, parent_idx);
-                rotate_right(forest, grandparent_idx);
+                // zig-zag:
+                rotate(forest, node_idx);
+                rotate(forest, node_idx);
             }
         } else {
             // zig
-            rotate_left(forest, parent_idx);
+            rotate(forest, node_idx);
         }
     }
 }
 
 #[cfg(test)]
 mod tests {
-    use super::{rotate_left, rotate_right};
+    use super::{rotate, rotate_left, rotate_right};
     use crate::node::{self, Node};
 
     fn create_nodes(n: usize) -> Vec<Node> {
@@ -165,7 +161,7 @@ mod tests {
 
     #[test]
     pub fn rotate_left_with_parent() {
-        // connect two nodes and rotate left on '0':
+        // form the tree and rotate left on '0':
         //      0                  2
         //     / \      =>        /
         //    1   2              0
@@ -201,7 +197,7 @@ mod tests {
 
     #[test]
     pub fn rotate_right_with_parent() {
-        // connect two nodes and rotate left on '0':
+        // form the tree and rotate left on '0':
         //      0               1
         //     / \      =>       \
         //    1   2               0
@@ -224,6 +220,42 @@ mod tests {
         assert!(forest[2].right.is_none());
     }
 
+    #[test]
+    pub fn rotate_parent_left() {
+        // form the tree and rotate on '1':
+        //      0              1
+        //     /       =>       \
+        //    1                  0
+        let mut forest = create_nodes(2);
+        forest[0].left = Some(1);
+        forest[1].parent = node::Parent::Node(0);
+        rotate(&mut forest, 1);
+        assert!(matches!(forest[1].parent, node::Parent::Root));
+        assert_eq!(forest[1].left, None);
+        assert_eq!(forest[1].right, Some(0));
+        assert!(matches!(forest[0].parent, node::Parent::Node(1)));
+        assert!(forest[0].left.is_none());
+        assert!(forest[0].right.is_none());
+    }
+
+    #[test]
+    pub fn rotate_parent_right() {
+        // form the tree and rotate on '1':
+        //    0                 1
+        //     \        =>     /
+        //      1             0
+        let mut forest = create_nodes(2);
+        forest[0].right = Some(1);
+        forest[1].parent = node::Parent::Node(0);
+        rotate(&mut forest, 1);
+        assert!(matches!(forest[1].parent, node::Parent::Root));
+        assert_eq!(forest[1].left, Some(0));
+        assert_eq!(forest[1].right, None);
+        assert!(matches!(forest[0].parent, node::Parent::Node(1)));
+        assert!(forest[0].left.is_none());
+        assert!(forest[0].right.is_none());
+    }
+
     #[test]
     pub fn splay_single_node() {
         // splay a single node, should do nothing: