diff --git a/src/lib.rs b/src/lib.rs index f3b3ea3..8eeb21e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -158,7 +158,9 @@ impl Node { (None, None) => *root = None, (Some(_), None) => *root = node.left.take(), (None, Some(_)) => *root = node.right.take(), - (Some(_), Some(_)) => node.value = Node::extract_min(&mut node.right), + (Some(_), Some(_)) => { + node.value = Node::remove_min(&mut node.right).unwrap() + } } Ok(()) @@ -183,13 +185,23 @@ impl Node { } } - fn extract_min(root: &mut HeapNode) -> T { + fn remove_min(root: &mut HeapNode) -> Option { if root.as_ref().unwrap().left.is_some() { - Node::extract_min(&mut root.as_mut().unwrap().left) + Node::remove_min(&mut root.as_mut().unwrap().left) } else { let node = root.take().unwrap(); *root = node.right; - node.value + Some(node.value) + } + } + + fn remove_max(root: &mut HeapNode) -> Option { + if root.as_ref().unwrap().right.is_some() { + Node::remove_max(&mut root.as_mut().unwrap().right) + } else { + let node = root.take().unwrap(); + *root = node.left; + Some(node.value) } } @@ -320,6 +332,32 @@ impl BinarySearchTree { } } + pub fn remove_min(&mut self) -> Option { + let removed_min = match self.root { + None => None, + Some(_) => Node::remove_min(&mut self.root), + }; + + if removed_min.is_some() { + self.size -= 1; + } + + removed_min + } + + pub fn remove_max(&mut self) -> Option { + let removed_max = match self.root { + None => None, + Some(_) => Node::remove_max(&mut self.root), + }; + + if removed_max.is_some() { + self.size -= 1; + } + + removed_max + } + pub fn pre_order(&self) -> Vec<&T> { let mut elements: Vec<&T> = Vec::new(); Node::pre_order_vec(&self.root, &mut elements); @@ -580,6 +618,44 @@ mod bst_test { assert_eq!(bst.max(), Some(&15)); } + #[test] + fn remove_min_from_bst() { + let mut bst = BinarySearchTree::empty(); + assert_eq!(bst.remove_min(), None); + + bst.insert(5); + assert_eq!(bst.remove_min(), Some(5)); + assert_eq!(bst.size(), 0); + + bst.insert(3); + bst.insert(1); + bst.insert(2); + bst.insert(15); + + assert_eq!(bst.remove_min(), Some(1)); + assert!(bst.contains(&2)); + assert_eq!(bst.size(), 3); + } + + #[test] + fn remove_max_from_bst() { + let mut bst = BinarySearchTree::empty(); + assert_eq!(bst.remove_max(), None); + + bst.insert(5); + assert_eq!(bst.remove_max(), Some(5)); + assert_eq!(bst.size(), 0); + + bst.insert(3); + bst.insert(1); + bst.insert(15); + bst.insert(10); + + assert_eq!(bst.remove_max(), Some(15)); + assert!(bst.contains(&10)); + assert_eq!(bst.size(), 3); + } + #[test] fn pre_order_traversal() { let mut bst = BinarySearchTree::empty();