From 39921553799cb7ebd3a68eda713088cc629b51e0 Mon Sep 17 00:00:00 2001 From: husky Date: Sun, 10 Sep 2023 15:42:31 -0700 Subject: [PATCH] add the ability to remove keys --- src/lib.rs | 341 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 341 insertions(+) diff --git a/src/lib.rs b/src/lib.rs index a94ca7a..137481b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,6 +2,7 @@ extern crate alloc; +use alloc::collections::VecDeque; use alloc::vec; use alloc::vec::Vec; @@ -150,6 +151,18 @@ impl BTree { } } + /// deletes a node from the internal node list, and updates all references to the following nodes so that they point to the correct node + pub fn delete_from_internal_list(&mut self, n: usize) { + for node in &mut self.nodes { + for child in &mut node.children { + if *child >= n && *child > 0 { + *child -= 1; + } + } + } + self.nodes.remove(n); + } + /// finds the node with the given key, or None if not found pub fn search(&mut self, key: u32) -> Option<&mut BTreeNode> { let mut index = self.root?; @@ -277,6 +290,291 @@ impl BTree { } } + /// removes the given key from this b-tree (THIS IS THE FUNCTION YOU LIKELY WANT) + pub fn remove(&mut self, key: u32) { + if let Some(mut root) = self.root { + let mut res = self.remove_from_nodetree(root, key); + let mut root_removed = false; + while let Some(node) = res.pop_front() { + if root >= node { + root -= 1; + } + if node == root { + root_removed = true; + } + } + self.root = Some(root); + // if root has no keys, make the first child the new root (if it has a child), else set root to None + if self.nodes[root].n == 0 || root_removed { + if self.nodes[root].leaf { + self.delete_from_internal_list(root); + self.root = None; + } else { + self.root = Some(self.nodes[root].children[0]); + self.delete_from_internal_list(root); + } + } + } + } + + /// removes the given key from the tree starting from a given node, returns VecDeque where Vec is a vector of nodes that were removed (so that other references can be updated) + pub fn remove_from_nodetree(&mut self, mut node: usize, key: u32) -> VecDeque { + let index = self.find_gte_key_on_node(node, key); + + let mut result = VecDeque::new(); + + if index < self.nodes[node].n as usize && self.nodes[node].keys[index].0 == key { + // if this is a leaf node, call remove_assuming_leaf + // vice versa + if self.nodes[node].leaf { + self.remove_assuming_leaf(node, index); + } else { + let res = self.remove_assuming_non_leaf(node, index); + result.append(&mut res.clone()); + } + + result + } else { + // if this is a leaf node, key not found + if self.nodes[node].leaf { + return result; + } + + // saved for later + let flag = index == self.nodes[node].n as usize; + + // if the child that should contain the key has less than m keys, steal from it + if self.nodes[self.nodes[node].children[index]].n < self.m { + let shift = self.steal_give(node, index); + if let Some(shift) = shift { + // if node is greater than shift, decrease them by 1 + if node >= shift { + node -= 1; + } + result.push_front(shift); + } + } + + // if the last child was merged, recurse on the (index-1)th child + // vice versa + if flag && index > self.nodes[node].n as usize { + let res = self.remove_from_nodetree(self.nodes[node].children[index - 1], key); + result.append(&mut res.clone()); + } else { + let res = self.remove_from_nodetree(self.nodes[node].children[index], key); + result.append(&mut res.clone()); + } + + result + } + } + + /// finds the first key on a node that is greater than or equal to the given key + pub fn find_gte_key_on_node(&mut self, node: usize, key: u32) -> usize { + let mut i = 0; + while i < self.nodes[node].n && self.nodes[node].keys[i as usize].0 < key { + i += 1; + } + i as usize + } + + /// steals keys from a node's children to fill child n, returns the position of the node that was removed so that other references can be updated, or None if no node was removed + pub fn steal_give(&mut self, node: usize, n: usize) -> Option { + // if the previous child has more than m-1 keys, steal from it + // else if the next child has more than m-1 keys, steal from it + // else, merge with one of the children + if n != 0 && self.nodes[self.nodes[node].children[n-1]].n >= self.m { + self.cfptn(node, n); + None + } else if n != self.nodes[node].n as usize && self.nodes[self.nodes[node].children[n+1]].n >= self.m { + self.cfftn(node, n); + None + } else if n != self.nodes[node].n as usize { + Some(self.merge_children(node, n).unwrap()) + } else { + Some(self.merge_children(node, n-1).unwrap()) + } + } + + /// copies the last key from the (n+1)th child and gives it to the nth child + pub fn cfftn(&mut self, node: usize, n: usize) { + let child = self.nodes[node].children[n]; + let sibling = self.nodes[node].children[n+1]; + + // give child the nth key from the current node + let child_n = self.nodes[child].n; + self.nodes[child].keys[child_n as usize] = self.nodes[node].keys[n].clone(); + + // if child isn't leaf, move sibling's first child as child's last child + if !self.nodes[child].leaf { + self.nodes[child].children[child_n as usize + 1] = self.nodes[sibling].children[0]; + } + + // move the first key from sibling to the current node + self.nodes[node].keys[n] = self.nodes[sibling].keys[0].clone(); + + // shift sibling keys to the left + for i in 1..self.nodes[sibling].n { + self.nodes[sibling].keys[i as usize - 1] = self.nodes[sibling].keys[i as usize].clone(); + } + + // if sibling isn't leaf, shift sibling children to the left + if !self.nodes[sibling].leaf { + for i in 1..self.nodes[sibling].n+1 { + self.nodes[sibling].children[i as usize - 1] = self.nodes[sibling].children[i as usize]; + } + } + + // decrease number of keys in sibling + self.nodes[sibling].n -= 1; + // increase number of keys in child + self.nodes[child].n += 1; + } + + /// copies the last key from the (n-1)th child and gives it to the nth child + pub fn cfptn(&mut self, node: usize, n: usize) { + let child = self.nodes[node].children[n]; + let sibling = self.nodes[node].children[n-1]; + + // move all keys in child one to the right + for i in (0..=self.nodes[child].n-1).rev() { + self.nodes[child].keys[i as usize + 1] = self.nodes[child].keys[i as usize].clone(); + } + + // if child isn't leaf, move all child pointers one step ahead + if !self.nodes[child].leaf { + for i in (0..=self.nodes[child].n).rev() { + self.nodes[child].children[i as usize + 1] = self.nodes[child].children[i as usize]; + } + } + + // give child the (n-1)th key from the current node + self.nodes[child].keys[0] = self.nodes[node].keys[n-1].clone(); + + // if child isn't leaf, move sibling's last child as child's first child + if !self.nodes[child].leaf { + self.nodes[child].children[0] = self.nodes[sibling].children[self.nodes[sibling].n as usize]; + } + + // move the (n-1)th key from sibling to the current node + self.nodes[node].keys[n-1] = self.nodes[sibling].keys[self.nodes[sibling].n as usize - 1].clone(); + + // decrease number of keys in sibling + self.nodes[sibling].n -= 1; + // increase number of keys in child + self.nodes[child].n += 1; + } + + /// merges the given node's nth child with its (n+1)th child, returns Some(n) if the merge was successful, None if the merge was not successful + /// where `n` is the index of the node that was removed so that other references can be updated + pub fn merge_children(&mut self, node: usize, n: usize) -> Option { + let child = self.nodes[node].children[n]; + let sibling = self.nodes[node].children[n+1]; + + // pop key from current node and insert it into child + self.nodes[child].keys[self.m as usize - 1] = self.nodes[node].keys[n].clone(); + + // copy keys from sibling into child + for i in 0..self.nodes[sibling].n { + self.nodes[child].keys[(self.m + i) as usize] = self.nodes[sibling].keys[i as usize].clone(); + } + + // if child isn't leaf, copy children from sibling into child + if !self.nodes[child].leaf { + for i in 0..self.nodes[sibling].n { + self.nodes[child].children[(self.m + i) as usize] = self.nodes[sibling].children[i as usize]; + } + } + + // move keys after n one to the left, filling the gap + for i in n+1..self.nodes[node].n as usize { + self.nodes[node].keys[i-1] = self.nodes[node].keys[i].clone(); + } + + // move children after n+1 one to the left, filling the gap + for i in n+2..self.nodes[node].n as usize + 1 { + self.nodes[node].children[i-1] = self.nodes[node].children[i]; + } + // decrease number of keys in node + self.nodes[child].n += self.nodes[sibling].n + 1; + self.nodes[node].n -= 1; + + // delete sibling + self.delete_from_internal_list(sibling); + + Some(sibling) + } + + /// removes the nth key from the given node assuming it is a non-leaf node, returns VecDeque where Vec is a vector of nodes that were removed (so that other references can be updated) + pub fn remove_assuming_non_leaf(&mut self, mut node: usize, n: usize) -> VecDeque { + let key = self.nodes[node].keys[n].0; + + let mut result = VecDeque::new(); + + let child_n = self.nodes[self.nodes[node].children[n]].n; + + // if child before k (children[n]) has at least m keys, find predecessor, replace key with predecessor, delete predecessor + // else if nth child has less than m keys, and (n+1)th child has at least m keys, replace key with successor, delete successor + // else merge nth child with (n+1)th child, recurse on nth child + if child_n >= self.m { + let pre_node = self.fkfl(node, n).unwrap().clone(); + self.nodes[node].keys[n] = pre_node.clone(); + let res = self.remove_from_nodetree(self.nodes[node].children[n], pre_node.0); + result.append(&mut res.clone()); + } else if self.nodes[self.nodes[node].children[n+1]].n >= self.m { + let suc_node = self.fkfr(node, n).unwrap().clone(); + self.nodes[node].keys[n] = suc_node.clone(); + let res = self.remove_from_nodetree(self.nodes[node].children[n+1], suc_node.0); + result.append(&mut res.clone()); + } else { + let res = self.merge_children(node, n).unwrap(); + if node >= res { + node -= 1; + result.push_front(res); + } + let res = self.remove_from_nodetree(self.nodes[node].children[n], key); + result.append(&mut res.clone()); + } + + result + } + + /// removes the nth key from the given node assuming it is a leaf node + pub fn remove_assuming_leaf(&mut self, node: usize, n: usize) { + // move all keys after n one to the left + for i in n+1..self.nodes[node].n as usize { + self.nodes[node].keys[i-1] = self.nodes[node].keys[i].clone(); + } + self.nodes[node].n -= 1; + } + + /// returns the last key of the last leaf node starting from the right of the children of the given node + pub fn fkfr(&self, node: usize, i: usize) -> Option<&(u32, T)> { + let children = &self.nodes[node].children; + if children.is_empty() { + return None; + } + let mut node = children[i]; + while !self.nodes[node].leaf { + node = self.nodes[node].children[self.nodes[node].n as usize]; + } + Some(&self.nodes[node].keys[self.nodes[node].n as usize - 1]) + } + + /// returns the first key of the first leaf node starting from the left of the children of the given node + pub fn fkfl(&self, node: usize, i: usize) -> Option<&(u32, T)> { + let children = &self.nodes[node].children; + if children.is_empty() { + return None; + } + let mut node = children[i]; + while !self.nodes[node].leaf { + node = self.nodes[node].children[0]; + } + Some(&self.nodes[node].keys[0]) + + } + /// traverses the tree in order and returns a vector of all keys pub fn traverse_in_order(&self, node: usize) -> Vec { let mut result = vec![]; @@ -293,6 +591,23 @@ impl BTree { result } + + /// traverses the tree in order and returns a vector of all values { + pub fn traverse_in_order_values(&self, node: usize) -> Vec { + let mut result = vec![]; + for i in 0..self.nodes[node].n { + if !self.nodes[node].leaf { + result.append(&mut self.traverse_in_order_values(self.nodes[node].children[i as usize])); + } + result.push(self.nodes[node].keys[i as usize].1.clone()); + } + + if !self.nodes[node].leaf { + result.append(&mut self.traverse_in_order_values(self.nodes[node].children[self.nodes[node].n as usize])); + } + + result + } } impl ToBytes for BTree { @@ -368,6 +683,11 @@ mod tests { tree.insert(3, 0); tree.insert(2, 0); tree.insert(1, 0); + tree.insert(24, 1); + tree.insert(25, 0); + + tree.remove(25); + tree.remove(24); fn tree_good(mut tree: BTree) { // print traversed tree @@ -385,5 +705,26 @@ mod tests { let tree2 = BTree::::from_bytes(&bytes); tree_good(tree2.clone()); + + // remove 10 and 20, then check if tree is still good + tree.remove(10); + assert_eq!(tree.traverse_in_order(tree.root.unwrap()), vec![1, 2, 3, 5, 6, 7, 11, 12, 17, 20, 30]); + tree.remove(20); + assert_eq!(tree.traverse_in_order(tree.root.unwrap()), vec![1, 2, 3, 5, 6, 7, 11, 12, 17, 30]); + + // remove all keys but 30 + let keys_remaining = vec![1, 2, 3, 5, 6, 7, 11, 12, 17]; + for key in keys_remaining { + tree.remove(key); + } + + // check if tree is still good + assert_eq!(tree.traverse_in_order(tree.root.unwrap()), vec![30]); + + // remove 30 + tree.remove(30); + + assert_eq!(tree.root, None); + assert_eq!(tree.nodes.len(), 0); } }