diff --git a/Cargo.toml b/Cargo.toml index 9a6cdc6..0c4332c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,3 +6,6 @@ edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] + +[dev-dependencies] +rand = "0.8" diff --git a/src/access.rs b/src/access.rs deleted file mode 100644 index 5575e41..0000000 --- a/src/access.rs +++ /dev/null @@ -1,71 +0,0 @@ -use crate::{ - node::{Node, Parent}, - splay::splay, -}; - -// constructs a path from a node to the root of the tree -pub fn access(forest: &mut Vec, node_idx: usize) { - assert!(node_idx < forest.len(), "access: node_idx out of bounds"); - splay(forest, node_idx); - - if let Some(right_idx) = forest[node_idx].right { - forest[node_idx].right = None; - forest[right_idx].parent = Parent::Path(node_idx); - } - - while let Parent::Path(path_idx) = forest[node_idx].parent { - splay(forest, path_idx); - // detach the right child of the path parent - if let Some(right_idx) = forest[path_idx].right { - forest[right_idx].parent = Parent::Path(path_idx); - forest[path_idx].right = None; - } - - // attach the node as the path parent's right child - forest[path_idx].right = Some(node_idx); - forest[node_idx].parent = Parent::Node(path_idx); - splay(forest, node_idx); - } -} - -#[cfg(test)] -mod tests { - use crate::node::{Node, Parent}; - - fn create_nodes(n: usize) -> Vec { - (0..n).map(|i| Node::new(i, 0.0)).collect() - } - - #[test] - pub fn base_case() { - // access a single node, should do nothing - let mut forest = create_nodes(1); - super::access(&mut forest, 0); - assert!(matches!(forest[0].parent, Parent::Root)); - } - - #[test] - pub fn access_leaf() { - let mut forest = create_nodes(3); - // '1' has a path pointer to '0', '1' has a right child '2'. - // after access(2), '2' should be the root of the tree: - // 0 0 0 2 - // | | \ / - // 1 => 2 => 2 => 0 - // \ / / \ - // 2 1 1 1 - forest[1].parent = Parent::Path(0); - forest[1].right = Some(2); - forest[2].parent = Parent::Node(1); - super::access(&mut forest, 2); - assert!(matches!(forest[2].parent, Parent::Root)); - assert_eq!(forest[2].right, None); - assert_eq!(forest[2].left, Some(0)); - assert!(matches!(forest[0].parent, Parent::Node(2))); - assert_eq!(forest[0].left, None); - assert_eq!(forest[0].right, Some(1)); - assert!(matches!(forest[1].parent, Parent::Node(0))); - assert_eq!(forest[1].left, None); - assert_eq!(forest[1].right, None); - } -} diff --git a/src/connected.rs b/src/connected.rs deleted file mode 100644 index 075f387..0000000 --- a/src/connected.rs +++ /dev/null @@ -1,72 +0,0 @@ -use crate::{ - access::access, - node::{Node, Parent}, -}; - -// checks if v and w are connected in the forest -pub fn connected(forest: &mut Vec, v: usize, w: usize) -> bool { - assert!( - v < forest.len() || w < forest.len(), - "connected: node_idx out of bounds" - ); - access(forest, v); // v is now the root of the tree - access(forest, w); - // if access(w) messed with the root of the tree, then v and w are connected: - !matches!(forest[v].parent, Parent::Root) || v == w -} - -#[cfg(test)] -mod tests { - use crate::node::{Node, Parent}; - - fn create_nodes(n: usize) -> Vec { - (0..n).map(|i| Node::new(i, 0.0)).collect() - } - - #[test] - pub fn base_case() { - // check two nodes, should return false - let mut forest = create_nodes(2); - assert!(!super::connected(&mut forest, 0, 1)); - } - - #[test] - pub fn connected_with_root() { - // check three nodes, one is root: - // 0 - // / \ - // 1 2 - let mut forest = create_nodes(3); - forest[0].left = Some(1); - forest[0].right = Some(2); - forest[1].parent = Parent::Node(0); - forest[2].parent = Parent::Node(0); - - assert!(super::connected(&mut forest, 0, 1)); - assert!(super::connected(&mut forest, 0, 2)); - assert!(super::connected(&mut forest, 1, 2)); - assert!(super::connected(&mut forest, 0, 0)); - assert!(super::connected(&mut forest, 1, 1)); - assert!(super::connected(&mut forest, 2, 2)); - } - - #[test] - pub fn connected_with_path_pointers() { - // check two trees that are connected by a path pointer - // 0 - // / \ - // 1 2 - // | - // 3 - let mut forest = create_nodes(4); - forest[0].left = Some(1); - forest[0].right = Some(2); - forest[1].parent = Parent::Node(0); - forest[2].parent = Parent::Node(0); - forest[3].parent = Parent::Path(2); - - assert!(super::connected(&mut forest, 0, 3)); - assert!(super::connected(&mut forest, 1, 3)); - assert!(super::connected(&mut forest, 2, 3)); - } -} diff --git a/src/cut.rs b/src/cut.rs deleted file mode 100644 index 61260ec..0000000 --- a/src/cut.rs +++ /dev/null @@ -1,96 +0,0 @@ -use crate::{ - access::access, - node::{Node, Parent}, -}; - -// cuts the link between v and its parent. -pub fn cut(forest: &mut Vec, v: usize) { - assert!(v < forest.len(), "cut: node_idx out of bounds"); - self::access(forest, v); - if let Some(left) = forest[v].left { - forest[left].parent = Parent::Root; - forest[v].left = None; - } -} - -#[cfg(test)] -mod tests { - use crate::{ - connected::connected, - cut::cut, - link::link, - node::{Node, Parent}, - }; - - fn create_nodes(n: usize) -> Vec { - (0..n).map(|i| Node::new(i, 0.0)).collect() - } - - #[test] - pub fn base_case() { - let mut forest = create_nodes(2); - assert!(!connected(&mut forest, 0, 1)); // not connected yet - - link(&mut forest, 0, 1); - assert!(matches!(forest[1].parent, Parent::Root)); - assert_eq!(forest[1].right, Some(0)); - assert_eq!(forest[1].left, None); - assert!(matches!(forest[0].parent, Parent::Node(1))); - // 1 - // \ <= link(0, 1) - // 0 - - assert!(connected(&mut forest, 0, 1)); // now connected - assert!(matches!(forest[1].parent, Parent::Root)); - assert_eq!(forest[1].right, None); - assert_eq!(forest[1].left, None); - assert!(matches!(forest[0].parent, Parent::Path(1))); - // 0 1 - // / => | - // 1 0 - - super::cut(&mut forest, 0); - assert!(matches!(forest[0].parent, Parent::Root)); - assert_eq!(forest[0].right, None); - // 0 0 - // / => - // 1 1 - assert!(!connected(&mut forest, 0, 1)); // now disconnected - } - - #[test] - pub fn cut_into_two_subtrees() { - let mut forest = create_nodes(5); - forest[0].left = Some(1); - forest[1].parent = Parent::Node(0); - forest[1].left = Some(2); - forest[2].parent = Parent::Node(1); - forest[3].right = Some(4); - forest[4].parent = Parent::Node(3); - // Given two trees: - // 0 3 - // / \ - // 1 4 - // / - // 2 - link(&mut forest, 2, 3); - // link(2, 3) should result in: - // 3 - // | \ - // 4 2 - // | - // 1 - // \ - // 0 - assert!(matches!(forest[3].parent, Parent::Root)); - assert!(matches!(forest[4].parent, Parent::Path(3))); - assert_eq!(forest[3].right, Some(2)); - assert!(matches!(forest[2].parent, Parent::Node(3))); - assert!(matches!(forest[1].parent, Parent::Path(2))); - assert_eq!(forest[1].right, Some(0)); - assert!(connected(&mut forest, 2, 3)); - cut(&mut forest, 2); - assert!(!connected(&mut forest, 2, 3)); - assert!(!connected(&mut forest, 2, 4)); - } -} diff --git a/src/lib.rs b/src/lib.rs index 9aba896..dc50993 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,7 +1,491 @@ #![allow(dead_code, clippy::module_name_repetitions)] // yes, I want to name my structs with the same name as the file -mod access; -mod connected; -mod cut; -mod link; mod node; +mod path; mod splay; + +use splay::update_max; + +use crate::{ + node::{Node, Parent}, + splay::splay, +}; + +pub struct LinkCutTree { + forest: Vec, +} + +impl LinkCutTree { + #[must_use] + pub fn new(n: usize) -> Self { + let nodes = (0..n).map(|i| Node::new(i, 0.0)).collect(); + Self { forest: nodes } + } + + // Constructs a path from a node to the root of the tree. + pub fn access(&mut self, v: usize) { + splay(&mut self.forest, v); + + if let Some(right_idx) = self.forest[v].right { + self.forest[v].right = None; + self.forest[right_idx].parent = Parent::Path(v); + } + + while let Parent::Path(path_idx) = self.forest[v].parent { + splay(&mut self.forest, path_idx); + // detach the right child of the path parent + if let Some(right_idx) = self.forest[path_idx].right { + self.forest[right_idx].parent = Parent::Path(path_idx); + self.forest[path_idx].right = None; + } + + // attach the node as the path parent's right child + self.forest[path_idx].right = Some(v); + self.forest[v].parent = Parent::Node(path_idx); + + splay(&mut self.forest, v); + } + + // update aggregate information + update_max(&mut self.forest, v); + } + + /// Creates a link between two nodes in the forest (v becomes the parent of w) + /// # Panics if v has a left child (i.e. v is not a root node) + pub fn link(&mut self, v: usize, w: usize) { + self.access(v); + self.access(w); + if !matches!(self.forest[v].parent, Parent::Root) || v == w { + return; // already connected + } + self.forest[w].right = Some(v); + self.forest[v].parent = Parent::Node(w); + } + + // Checks if v and w are connected in the forest + pub fn connected(&mut self, v: usize, w: usize) -> bool { + self.access(v); // v is now the root of the tree + self.access(w); + // if access(w) messed with the root of the tree, then v and w are connected: + !matches!(self.forest[v].parent, Parent::Root) || v == w + } + + // Cuts the link between v and its parent. + pub fn cut(&mut self, v: usize) { + self.access(v); + if let Some(left) = self.forest[v].left { + self.forest[left].parent = Parent::Root; + self.forest[v].left = None; + } + } + + // Finds the maximum weight in the path from v and its parent + pub fn findmax(&mut self, v: usize) -> usize { + self.access(v); + self.forest[v].max_weight_idx + } +} + +#[cfg(test)] +mod tests { + use crate::node::Parent; + use rand::{seq::SliceRandom, Rng, SeedableRng}; + + #[test] + pub fn access_base_case() { + // access a single node, should do nothing + let mut tree = super::LinkCutTree::new(1); + tree.access(0); + assert!(matches!(tree.forest[0].parent, Parent::Root)); + } + + #[test] + pub fn access_leaf() { + let mut tree = super::LinkCutTree::new(4); + // '1' has a path pointer to '0', '1' has a right child '2'. + // after access(2), '2' should be the root of the tree: + // 0 0 0 2 + // | | \ / + // 1 => 2 => 2 => 0 + // \ / / \ + // 2 1 1 1 + tree.forest[1].parent = Parent::Path(0); + tree.forest[1].right = Some(2); + tree.forest[2].parent = Parent::Node(1); + tree.access(2); + assert!(matches!(tree.forest[2].parent, Parent::Root)); + assert_eq!(tree.forest[2].right, None); + assert_eq!(tree.forest[2].left, Some(0)); + assert!(matches!(tree.forest[0].parent, Parent::Node(2))); + assert_eq!(tree.forest[0].left, None); + assert_eq!(tree.forest[0].right, Some(1)); + assert!(matches!(tree.forest[1].parent, Parent::Node(0))); + assert_eq!(tree.forest[1].left, None); + assert_eq!(tree.forest[1].right, None); + } + + #[test] + pub fn link_base_case() { + let mut tree = super::LinkCutTree::new(2); + assert!(!tree.connected(0, 1)); // not connected yet + tree.link(0, 1); + assert!(tree.connected(0, 1)); // now connected + } + + #[test] + pub fn link_already_connected() { + // '2' has a right child '3': + // link(0, 3) should add no link, and result in (| denotes a path pointer): + // 0 0 0 3 + // / \ / | / | / + // 1 2 => 1 2 => 1 3 => 0 + // \ \ / / \ + // 3 3 2 1 2 + // + let mut tree = super::LinkCutTree::new(4); + tree.forest[0].left = Some(1); + tree.forest[0].right = Some(2); + tree.forest[1].parent = Parent::Node(0); + tree.forest[2].parent = Parent::Node(0); + tree.forest[2].right = Some(3); + tree.forest[3].parent = Parent::Node(2); + tree.link(0, 3); + assert!(matches!(tree.forest[3].parent, Parent::Root)); + assert_eq!(tree.forest[3].left, Some(0)); + assert_eq!(tree.forest[3].right, None); + assert!(matches!(tree.forest[0].parent, Parent::Node(3))); + assert_eq!(tree.forest[0].left, Some(1)); + assert_eq!(tree.forest[0].right, Some(2)); + assert!(matches!(tree.forest[1].parent, Parent::Node(0))); + assert_eq!(tree.forest[1].left, None); + assert_eq!(tree.forest[1].right, None); + assert!(matches!(tree.forest[2].parent, Parent::Node(0))); + assert_eq!(tree.forest[2].left, None); + assert_eq!(tree.forest[2].right, None); + } + + #[test] + pub fn link_already_connected_with_path() { + // '3' has a path pointer to '2', and '2' has a path pointer to '0': + // link(0, 3) should add no link, and result in (| denotes a path pointer): + // 0 0 0 3 + // / \ / | / | / + // 1 2 => 1 2 => 1 3 => 0 + // | | / / \ + // 3 3 2 1 2 + // + let mut tree = super::LinkCutTree::new(4); + tree.forest[0].left = Some(1); + tree.forest[0].right = Some(2); + tree.forest[1].parent = Parent::Node(0); + tree.forest[2].parent = Parent::Node(0); + tree.forest[3].parent = Parent::Path(2); + tree.link(0, 3); + assert!(matches!(tree.forest[3].parent, Parent::Root)); + assert_eq!(tree.forest[3].left, Some(0)); + assert_eq!(tree.forest[3].right, None); + assert!(matches!(tree.forest[0].parent, Parent::Node(3))); + assert_eq!(tree.forest[0].left, Some(1)); + assert_eq!(tree.forest[0].right, Some(2)); + assert!(matches!(tree.forest[1].parent, Parent::Node(0))); + assert_eq!(tree.forest[1].left, None); + assert_eq!(tree.forest[1].right, None); + assert!(matches!(tree.forest[2].parent, Parent::Node(0))); + assert_eq!(tree.forest[2].left, None); + assert_eq!(tree.forest[2].right, None); + } + + #[test] + pub fn link_to_leftmost() { + // Given two trees: + // 0 3 + // / \ + // 1 2 + // link(3, 1) should result in a single tree (| denotes a path pointer): + // 1 3 1 + // | |\ + // 0 => 0 3 + // \ \ + // 2 2 + let mut tree = super::LinkCutTree::new(4); + tree.forest[0].left = Some(1); + tree.forest[0].right = Some(2); + tree.forest[1].parent = Parent::Node(0); + tree.forest[2].parent = Parent::Node(0); + tree.link(3, 1); + assert!(matches!(tree.forest[1].parent, Parent::Root)); + assert_eq!(tree.forest[1].left, None); + assert_eq!(tree.forest[1].right, Some(3)); + assert!(matches!(tree.forest[0].parent, Parent::Path(1))); + assert_eq!(tree.forest[0].left, None); + assert_eq!(tree.forest[0].right, Some(2)); + } + + #[test] + pub fn connected_with_root() { + // check three nodes, one is root: + // 0 + // / \ + // 1 2 + let mut tree = super::LinkCutTree::new(3); + tree.forest[0].left = Some(1); + tree.forest[0].right = Some(2); + tree.forest[1].parent = Parent::Node(0); + tree.forest[2].parent = Parent::Node(0); + + assert!(tree.connected(0, 1)); + assert!(tree.connected(0, 2)); + assert!(tree.connected(1, 2)); + assert!(tree.connected(0, 0)); + assert!(tree.connected(1, 1)); + assert!(tree.connected(2, 2)); + } + + #[test] + pub fn connected_with_path_pointers() { + // check two splay trees that are connected by a path pointer + // 0 + // / \ + // 1 2 + // | + // 3 + let mut tree = super::LinkCutTree::new(4); + tree.forest[0].left = Some(1); + tree.forest[0].right = Some(2); + tree.forest[1].parent = Parent::Node(0); + tree.forest[2].parent = Parent::Node(0); + tree.forest[3].parent = Parent::Path(2); + + assert!(tree.connected(0, 3)); + assert!(tree.connected(1, 3)); + assert!(tree.connected(2, 3)); + } + + #[test] + pub fn cut_base_case() { + let mut tree = super::LinkCutTree::new(2); + assert!(!tree.connected(0, 1)); // not connected yet + + tree.link(0, 1); + // 1 + // \ <= link(0, 1) + // 0 + assert!(matches!(tree.forest[1].parent, Parent::Root)); + assert_eq!(tree.forest[1].left, None); + assert_eq!(tree.forest[1].right, Some(0)); + assert!(matches!(tree.forest[0].parent, Parent::Node(1))); + + assert!(tree.connected(0, 1)); // now connected + assert!(matches!(tree.forest[1].parent, Parent::Root)); + assert_eq!(tree.forest[1].right, None); + assert_eq!(tree.forest[1].left, None); + assert!(matches!(tree.forest[0].parent, Parent::Path(1))); + + tree.cut(0); + assert!(matches!(tree.forest[0].parent, Parent::Root)); + assert_eq!(tree.forest[0].right, None); + // 1 1 + // \ => + // 0 0 + assert!(!tree.connected(0, 1)); // now disconnected + } + + #[test] + pub fn cut_into_two_subtrees() { + let mut tree = super::LinkCutTree::new(5); + tree.forest[0].left = Some(1); + tree.forest[1].parent = Parent::Node(0); + tree.forest[1].left = Some(2); + tree.forest[2].parent = Parent::Node(1); + tree.forest[3].right = Some(4); + tree.forest[4].parent = Parent::Node(3); + // Given two trees: + // 0 3 + // / \ + // 1 4 + // / + // 2 + tree.link(2, 3); + // link(2, 3) should result in: + // 3 + // | \ + // 4 2 + // | + // 1 + // \ + // 0 + assert!(matches!(tree.forest[3].parent, Parent::Root)); + assert_eq!(tree.forest[3].left, None); + assert_eq!(tree.forest[3].right, Some(2)); + assert!(matches!(tree.forest[4].parent, Parent::Path(3))); + assert!(matches!(tree.forest[2].parent, Parent::Node(3))); + assert_eq!(tree.forest[2].left, None); + assert_eq!(tree.forest[2].right, None); + assert!(matches!(tree.forest[1].parent, Parent::Path(2))); + assert_eq!(tree.forest[1].left, None); + assert_eq!(tree.forest[1].right, Some(0)); + assert!(matches!(tree.forest[0].parent, Parent::Node(1))); + assert!(tree.connected(2, 3)); + + // we cut node 2 from its parent 3: + tree.cut(2); + assert!(!tree.connected(2, 3)); + assert!(!tree.connected(2, 4)); + } + + #[test] + pub fn connectivity() { + // We form a link-cut tree from a rooted tree with the following structure: + // 0 + // / \ + // 1 6 + // / \ \ + // 2 3 7 + // / \ \ + // 4 5 8 + // / + // 9 + let mut lctree = super::LinkCutTree::new(10); + lctree.link(1, 0); + lctree.link(2, 1); + lctree.link(3, 1); + lctree.link(4, 3); + lctree.link(5, 3); + lctree.link(6, 0); + lctree.link(7, 6); + lctree.link(8, 7); + lctree.link(9, 8); + + // checking connectivity: + for i in 0..10 { + for j in 0..10 { + assert!(lctree.connected(i, j)); + } + } + + // we cut node 6 from its parent 0: + lctree.cut(6); + + // the forest should now look like this: + // 0 + // / + // 1 6 + // / \ \ + // 2 3 7 + // / \ \ + // 4 5 8 + // / + // 9 + + // we check connectivity again for the two trees: + for i in 0..6 { + for j in 0..6 { + assert!(lctree.connected(i, j)); + } + } + for i in 6..10 { + for j in 6..10 { + assert!(lctree.connected(i, j)); + } + } + for i in 0..6 { + for j in 6..10 { + assert!(!lctree.connected(i, j)); + } + } + } + + #[test] + pub fn findmax() { + let mut lctree = super::LinkCutTree::new(7); + lctree.forest[0].weight = 5.0; + lctree.forest[1].weight = 2.0; + lctree.forest[2].weight = 6.0; + lctree.forest[3].weight = 3.0; + lctree.forest[4].weight = 4.0; + lctree.forest[5].weight = 0.0; + lctree.forest[6].weight = 1.0; + + // we form the following structure: + // 0(5) + // / + // 1(2) + // / + // 2(6) + // / + // 3(3) + // / \ + // 4(4) 6(1) + // / + // 5(0) + + lctree.link(1, 0); + lctree.link(2, 1); + lctree.link(3, 2); + lctree.link(4, 3); + lctree.link(5, 4); + lctree.link(6, 3); + + assert_eq!(lctree.findmax(2), 2); + assert_eq!(lctree.findmax(6), 2); + assert_eq!(lctree.findmax(0), 0); + assert_eq!(lctree.findmax(5), 2); + assert_eq!(lctree.findmax(3), 2); + assert_eq!(lctree.findmax(1), 0); + assert_eq!(lctree.findmax(4), 2); + + // we cut node 3 from its parent 2: + lctree.cut(3); + assert_eq!(lctree.findmax(5), 4); + assert_eq!(lctree.findmax(6), 3); + assert_eq!(lctree.findmax(4), 4); + assert_eq!(lctree.findmax(3), 3); + assert_eq!(lctree.findmax(2), 2); + assert_eq!(lctree.findmax(1), 0); + assert_eq!(lctree.findmax(0), 0); + } + + fn create_random_tree(n: usize, seed: u64) -> (Vec<(usize, usize)>, Vec, Vec) { + let mut rng = rand::rngs::StdRng::seed_from_u64(seed); + + let mut edges = Vec::new(); + let mut weights: Vec = (0..n).map(|i| i as f64).collect(); + weights.shuffle(&mut rng); + + let mut ground_truth = vec![0; n]; + let mut in_tree = Vec::from([0]); + for i in 1..n { + let parent_idx = rng.gen_range(0..in_tree.len()); + let parent = in_tree[parent_idx]; + edges.push((i, parent)); + + ground_truth[i] = if weights[i] > weights[ground_truth[parent]] { + i + } else { + ground_truth[parent] + }; + in_tree.push(i); + } + + (edges, weights, ground_truth) + } + + #[test] + pub fn findmax_random() { + let n = 100; + let seed = rand::thread_rng().gen(); + let (edges, weights, ground_truth) = create_random_tree(n, seed); + let mut lctree = super::LinkCutTree::new(n); + for i in 0..n { + lctree.forest[i].weight = weights[i]; + } + + for (v, w) in edges { + lctree.link(v, w); + } + + for _ in 0..n * 100 { + let v = rand::thread_rng().gen_range(0..n); + assert_eq!(lctree.findmax(v), ground_truth[v]); + } + } +} diff --git a/src/link.rs b/src/link.rs deleted file mode 100644 index 27e6377..0000000 --- a/src/link.rs +++ /dev/null @@ -1,137 +0,0 @@ -use crate::{ - access::access, - node::{Node, Parent}, -}; - -// creates a link between two nodes in the forest (w becomes the parent of v) -pub fn link(forest: &mut Vec, v: usize, w: usize) { - assert!( - v < forest.len() && w < forest.len(), - "link: node_idx out of bounds" - ); - access(forest, v); - access(forest, w); - if !matches!(forest[v].parent, Parent::Root) || v == w { - return; // already connected - } - assert!(forest[w].right.is_none(), "link: w should be a root!"); - forest[w].right = Some(v); - forest[v].parent = Parent::Node(w); -} - -#[cfg(test)] -mod tests { - use super::link; - use crate::{ - connected::connected, - node::{Node, Parent}, - }; - - fn create_nodes(n: usize) -> Vec { - (0..n).map(|i| Node::new(i, 0.0)).collect() - } - - #[test] - pub fn base_case() { - let mut forest = create_nodes(2); - assert!(!connected(&mut forest, 0, 1)); // not connected yet - super::link(&mut forest, 0, 1); - assert!(connected(&mut forest, 0, 1)); // now connected - } - - #[test] - pub fn already_connected() { - // '2' has a right child '3': - // link(0, 3) should add no link, and result in (| denotes a path pointer): - // 0 0 0 3 - // / \ / | / | / - // 1 2 => 1 2 => 1 3 => 0 - // \ \ / / \ - // 3 3 2 1 2 - // - let mut forest = create_nodes(4); - forest[0].left = Some(1); - forest[0].right = Some(2); - forest[1].parent = Parent::Node(0); - forest[2].parent = Parent::Node(0); - forest[2].right = Some(3); - forest[3].parent = Parent::Node(2); - link(&mut forest, 0, 3); - assert!(matches!(forest[3].parent, Parent::Root)); - assert_eq!(forest[3].left, Some(0)); - assert_eq!(forest[3].right, None); - assert!(matches!(forest[0].parent, Parent::Node(3))); - assert_eq!(forest[0].left, Some(1)); - assert_eq!(forest[0].right, Some(2)); - assert!(matches!(forest[1].parent, Parent::Node(0))); - assert_eq!(forest[1].left, None); - assert_eq!(forest[1].right, None); - assert!(matches!(forest[2].parent, Parent::Node(0))); - assert_eq!(forest[2].left, None); - assert_eq!(forest[2].right, None); - } - - #[test] - pub fn already_connected_with_path() { - // '3' has a path pointer to '2', and '2' has a path pointer to '0': - // link(0, 1) should add no link, and result in (| denotes a path pointer): - // 0 0 0 3 - // / \ / | / | / - // 1 2 => 1 2 => 1 3 => 0 - // | | / / \ - // 3 3 2 1 2 - // - let mut forest = create_nodes(4); - forest[0].left = Some(1); - forest[0].right = Some(2); - forest[1].parent = Parent::Node(0); - forest[2].parent = Parent::Node(0); - forest[3].parent = Parent::Path(2); - link(&mut forest, 0, 3); - assert!(matches!(forest[3].parent, Parent::Root)); - assert_eq!(forest[3].left, Some(0)); - assert_eq!(forest[3].right, None); - assert!(matches!(forest[0].parent, Parent::Node(3))); - assert_eq!(forest[0].left, Some(1)); - assert_eq!(forest[0].right, Some(2)); - assert!(matches!(forest[1].parent, Parent::Node(0))); - assert_eq!(forest[1].left, None); - assert_eq!(forest[1].right, None); - assert!(matches!(forest[2].parent, Parent::Node(0))); - assert_eq!(forest[2].left, None); - assert_eq!(forest[2].right, None); - } - - #[test] - pub fn link_to_leftmost() { - // Given two trees: - // 0 3 - // / \ - // 1 2 - // link(1, 3) should result in a single tree (| denotes a path pointer): - // 1 3 1 - // | | \ - // 0 => 0 3 - // \ \ - // 2 2 - // - let mut forest = create_nodes(4); - forest[0].left = Some(1); - forest[0].right = Some(2); - forest[1].parent = Parent::Node(0); - forest[2].parent = Parent::Node(0); - link(&mut forest, 3, 1); - assert!(matches!(forest[1].parent, Parent::Root)); - assert_eq!(forest[1].right, Some(3)); - assert_eq!(forest[1].left, None); - assert!(matches!(forest[3].parent, Parent::Node(1))); - assert_eq!(forest[3].right, None); - assert_eq!(forest[3].left, None); - assert!(matches!(forest[0].parent, Parent::Path(1))); - assert_eq!(forest[0].right, Some(2)); - assert_eq!(forest[0].left, None); - assert!(matches!(forest[2].parent, Parent::Node(0))); - assert_eq!(forest[2].right, None); - assert_eq!(forest[2].left, None); - } -} diff --git a/src/node.rs b/src/node.rs index 889a62f..6b3585c 100644 --- a/src/node.rs +++ b/src/node.rs @@ -5,46 +5,39 @@ pub enum Parent { Root, // root of the tree } -pub struct Path { - pub max_weight_idx: usize, - pub max_weight: f64, -} - -impl Path { - pub fn new(max_weight_idx: usize, max_weight: f64) -> Self { - Path { - max_weight_idx, - max_weight, - } - } - - pub fn update(&mut self, idx: usize, weight: f64) { - if weight > self.max_weight { - self.max_weight = weight; - self.max_weight_idx = idx; - } - } -} - pub struct Node { + pub idx: usize, pub left: Option, pub right: Option, pub parent: Parent, pub flipped: bool, - // path aggregates + // for path aggregation: pub weight: f64, - pub path: Path, + pub max_weight_idx: usize, } impl Node { pub fn new(idx: usize, weight: f64) -> Self { Node { + idx, left: None, right: None, parent: Parent::Root, flipped: false, - weight: 0.0, - path: Path::new(idx, weight), + weight, + max_weight_idx: idx, } } + + pub fn to_str(&self) -> String { + let parent = match self.parent { + Parent::Node(idx) => format!("Node({idx})"), + Parent::Path(idx) => format!("Path({idx})"), + Parent::Root => "Root".to_string(), + }; + format!( + "Node {{ idx: {}, left: {:?}, right: {:?}, parent: {:?}, max_weight_idx: {} }}", + self.idx, self.left, self.right, parent, self.max_weight_idx + ) + } } diff --git a/src/path.rs b/src/path.rs new file mode 100644 index 0000000..fd19a61 --- /dev/null +++ b/src/path.rs @@ -0,0 +1,43 @@ +use crate::node::Node; + +pub trait Path { + fn aggregate(&mut self, other: &Node); +} + +pub struct FindMax { + pub max_weight_idx: usize, + pub max_weight: f64, +} + +impl Path for FindMax { + fn aggregate(&mut self, other: &Node) { + if other.weight > self.max_weight { + self.max_weight = other.weight; + self.max_weight_idx = other.idx; + } + } +} + +pub struct FindMin { + pub min_weight_idx: usize, + pub min_weight: f64, +} + +impl Path for FindMin { + fn aggregate(&mut self, other: &Node) { + if other.weight < self.min_weight { + self.min_weight = other.weight; + self.min_weight_idx = other.idx; + } + } +} + +pub struct FindSum { + pub sum: f64, +} + +impl Path for FindSum { + fn aggregate(&mut self, other: &Node) { + self.sum += other.weight; + } +} diff --git a/src/splay.rs b/src/splay.rs index 5225824..e23d4cb 100644 --- a/src/splay.rs +++ b/src/splay.rs @@ -1,10 +1,10 @@ use crate::node::{Node, Parent}; -fn rotate_left(forest: &mut Vec, node_idx: usize) { +fn rotate_left(forest: &mut [Node], node_idx: usize) { // base cases: // - node_idx is out of bounds // - node_idx should have a right child - if node_idx >= forest.len() || forest[node_idx].right.is_none() { + if forest[node_idx].right.is_none() { return; } @@ -25,13 +25,17 @@ fn rotate_left(forest: &mut Vec, node_idx: usize) { if let Some(new_right_child) = forest[node_idx].right { forest[new_right_child].parent = Parent::Node(node_idx); } + + // update aggregate information + update_max(forest, node_idx); + update_max(forest, right_child); } -fn rotate_right(forest: &mut Vec, node_idx: usize) { +fn rotate_right(forest: &mut [Node], node_idx: usize) { // base cases: // - node_idx is out of bounds // - node_idx should have a left child - if node_idx >= forest.len() || forest[node_idx].left.is_none() { + if forest[node_idx].left.is_none() { return; } @@ -52,10 +56,31 @@ fn rotate_right(forest: &mut Vec, node_idx: usize) { if let Some(new_left_child) = forest[node_idx].left { forest[new_left_child].parent = Parent::Node(node_idx); } + + // update aggregate information + update_max(forest, node_idx); + 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; + } + } + forest[node_idx].max_weight_idx = max_idx; } // makes node_idx the root of its Splay tree -pub fn splay(forest: &mut Vec, node_idx: usize) { +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) {