Mercurial > lbo > hg > btreex
view src/lib.rs @ 1:04b6ed18abf3
Better splitting
author | Lewin Bormann <lbo@spheniscida.de> |
---|---|
date | Mon, 04 Jul 2022 18:14:18 -0700 |
parents | e342ded1647e |
children | df21e58157f3 |
line wrap: on
line source
use ptree::TreeBuilder; use std::cmp::Ordering; const ORDER: usize = 5; #[derive(Debug)] struct Node<T> { entries: [Option<T>; ORDER - 1], children: [Option<usize>; ORDER], parent: Option<usize>, } impl<T: Ord> Node<T> { fn new(parent: Option<usize>) -> Node<T> { Node { entries: Default::default(), children: Default::default(), parent: parent, } } } #[derive(Debug)] struct BTree<T> { v: Vec<Node<T>>, root: usize, } impl<T: std::fmt::Debug + Ord> BTree<T> { fn new(initial_capacity: usize) -> BTree<T> { BTree { v: Vec::with_capacity(initial_capacity), root: 0, } } fn free(&self, node: usize) -> usize { self.v[node] .entries .iter() .map(|e| if e.is_some() { 0 } else { 1 }) .sum() } fn get(&self, node: usize, elem: usize) -> &Option<T> { &self.v[node].entries[elem] } fn get_mut(&mut self, node: usize, elem: usize) -> &mut Option<T> { &mut self.v[node].entries[elem] } fn get_child(&self, node: usize, child: usize) -> &Option<usize> { &self.v[node].children[child] } fn get_child_mut(&mut self, node: usize, child: usize) -> &mut Option<usize> { &mut self.v[node].children[child] } fn insert_inline(&mut self, node: usize, value: T) { assert!(self.free(node) >= 1); let node = &mut self.v[node]; let mut loc = ORDER - 2; for i in 0..ORDER - 1 { if let Some(ref el) = node.entries[i] { if value.cmp(el) == Ordering::Less { // Found right location. Shift remaining elements. loc = i; break; } } else { loc = i; break; } } for i in 0..(ORDER - 2 - loc) { node.entries[ORDER - 2 - i] = node.entries[ORDER - 3 - i].take(); node.children[ORDER - 2 - i] = node.children[ORDER - 3 - i].take(); } node.entries[loc] = Some(value); } fn insert(&mut self, value: T) -> bool { // 1. Find appropriate node // 2. Find appropriate order // 3. Insert or split. // Initialize empty tree if self.v.is_empty() { self.v.push(Node::new(None)); self.v[0].entries[0] = Some(value); return true; } let mut node = self.root; 'CHILDREN: loop { // 1. Descend as long as there are feasible children. // 2. Insert or split once no more children to be found. let mut found = false; let mut i = 0; 'ELEMENTS: while i < ORDER - 1 { if let Some(el) = self.get(node, i) { if value.cmp(el) == Ordering::Less { found = true; break 'ELEMENTS; } } else { // Past last filled element. found = true; break 'ELEMENTS; } i += 1 } // There is an element greater than `value`? if let Some(child) = self.get_child(node, i) { // Descend into child node. node = *child; continue 'CHILDREN; } else if self.free(node) >= 1 { self.insert_inline(node, value); return true; } else { // We have to split. //self.split(node, Some(value), i); let new = self.split(node, None, 0); node = new; continue; } } } /// Split node `node` at `loc`, moving all values at `loc` and higher to the right node, /// inserting `value` into a new parent node. fn split(&mut self, node: usize, value: Option<T>, loc: usize) -> usize { let leftix = node; let rightix = self.v.len(); let topix = self.v.len() + 1; let mut right = Node::new(Some(topix)); let mut top = Node::new(self.v[node].parent); // TODO: We create a top node. In reality, the pivot value is to be inserted into the // existing parent node. if value.is_some() { // Put pivot entry into top node. top.entries[0] = value; // Copy greater elements to right node. for i in loc..ORDER - 1 { right.entries[i - loc] = self.get_mut(node, i).take(); } // left top.children[0] = Some(leftix); // right top.children[1] = Some(rightix); } else { // Split without inserting value. let pivotix = (ORDER-1) / 2; let pivot = self.get_mut(node, pivotix).take(); top.entries[0] = pivot; // Copy greater elements to right node. for i in pivotix..ORDER - 2 { right.entries[i - pivotix] = self.get_mut(node, i+1).take(); println!("Inserting {:?} at {}", right.entries[i-pivotix], i-pivotix); } top.children[0] = Some(leftix); top.children[1] = Some(rightix); } // Update parent. { if let Some(parentix) = self.v[leftix].parent { let parent = &mut self.v[parentix]; for i in 0..ORDER { if parent.children[i] == Some(node) { parent.children[i] = Some(topix); break; } } } else { self.root = topix; } } self.v[leftix].parent = Some(topix); right.parent = Some(topix); self.v.push(right); self.v.push(top); topix } } impl<T: std::fmt::Debug> BTree<T> { fn debug_print(&self) -> String { let mut t = TreeBuilder::new("btree".to_string()); self.print_subtree(&mut t, self.root); let mut o: Vec<u8> = Vec::new(); ptree::write_tree(&t.build(), &mut o); String::from_utf8(o).unwrap() } fn print_subtree(&self, tb: &mut TreeBuilder, node: usize) { let node = &self.v[node]; for i in 0..ORDER - 1 { if let Some(ch) = node.children[i] { self.print_subtree(tb.begin_child(format!("<{}>", ch)), ch); tb.end_child(); } tb.add_empty_child(format!("{:?}", node.entries[i])); } } } #[cfg(test)] mod tests { use super::BTree; #[test] fn basic_test() { let mut bt = BTree::<i32>::new(16); let ns = &[2,4,6,8,10,12,14,1,16,18,20,7,22,24,26,28,30]; for n in ns { bt.insert(*n); println!("Inserting {}", n); println!("root: {}", bt.root); println!("{}", bt.debug_print()); } for (i, e) in bt.v.iter().enumerate() { println!("{} {:?}", i, e); } } }