From 09bfb91365c57a097ac94e0029e03370dfd52710 Mon Sep 17 00:00:00 2001 From: Aziz Kayumov Date: Fri, 6 Oct 2023 15:08:34 +0900 Subject: [PATCH 1/2] clean up splay --- src/lib.rs | 4 +- src/path.rs | 17 ++++++++ src/splay.rs | 113 ++++++++++++++++++++++++++++++++------------------- 3 files changed, 92 insertions(+), 42 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 404bdbe..12f0633 100644 --- 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 index fd19a61..33da558 100644 --- 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 index 1827e8f..dd8025b 100644 --- a/src/splay.rs +++ b/src/splay.rs @@ -1,4 +1,4 @@ -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 +84,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. @@ -114,41 +122,28 @@ 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 { @@ -165,7 +160,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 +196,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 +219,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: From 0e108e57b94cafac4155a1bd141febe1fcef8c32 Mon Sep 17 00:00:00 2001 From: Aziz Kayumov Date: Fri, 6 Oct 2023 16:37:09 +0900 Subject: [PATCH 2/2] make clippy happy --- src/splay.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/splay.rs b/src/splay.rs index dd8025b..761b034 100644 --- a/src/splay.rs +++ b/src/splay.rs @@ -1,4 +1,7 @@ -use crate::{node::{Node, Parent}, path::update_max}; +use crate::{ + node::{Node, Parent}, + path::update_max, +}; /// Rotates the subtree rooted at `node_idx` to the left: /// @@ -119,8 +122,6 @@ fn rotate(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 let Parent::Node(grandparent_idx) = forest[parent_idx].parent { if (forest[grandparent_idx].left == Some(parent_idx))