use crate::{Node, Side, ARRAY_SIZE};
use std::cmp::{
max,
Ordering::{self, Equal, Greater, Less},
};
impl<K: Default + Clone, V: Default + Clone> Node<K, V> {
pub fn is_balanced(&mut self) -> bool {
self.update_child_difference();
self.child_difference < 2
}
pub fn balance(&mut self) {
self.diff_balance();
self.height_balance();
}
pub fn diff_balance(&mut self) {
if !self.is_balanced() {
let left_child = match &self.children[0] {
None => 0,
Some(child) => child.child_difference,
};
let right_child = match &self.children[1] {
None => 0,
Some(child) => child.child_difference,
};
let side: Option<Side> = match left_child.cmp(&right_child) {
Less => Some(Side::Left),
Equal => None,
Greater => Some(Side::Right),
};
match side {
None => {}
Some(side) => self.single_rotation(side),
}
}
}
pub fn has_full_child(&self, index: (usize, usize)) -> (bool, bool) {
match &self.children[index.1] {
None => (false, false),
Some(child) => (child.children[index.0].is_some(), child.is_full()),
}
}
pub fn single_rotation(&mut self, side: Side) {
let index = match side {
Side::Left => (0, 1),
Side::Right => (1, 0),
};
self.children[index.0] = Some({
let mut child = Node::new();
child.items = self.take_items();
child.children[index.0] = self.take_child(index.0);
child.children_size[index.0] = self.take_child_size(index.0);
child.update_child_difference();
child
});
self.children_size[index.0] += ARRAY_SIZE;
self.children_size[index.1] -= ARRAY_SIZE;
match &self.has_full_child(index) {
(false, _) => self.child_to_parent(index),
(true, false) => {
self.child_to_parent(index);
self.children[index.0].as_deref_mut().unwrap().children[index.1] =
match self.children[index.1].as_deref_mut() {
None => None,
Some(child) => child.take_child(index.0),
};
self.children_size[index.0] += match self.children[index.1].as_deref_mut() {
None => 0,
Some(child) => child.take_child_size(index.0),
};
self.update_child_difference()
}
(true, true) => {
self.subchild_to_parent(index);
}
}
}
pub fn child_to_parent(&mut self, index: (usize, usize)) {
self.items = self.children[index.1].as_deref_mut().unwrap().take_items();
self.children[index.1] = self.children[index.1]
.as_deref_mut()
.unwrap()
.take_child(index.1);
match self.children[index.1].as_deref_mut() {
None => {}
Some(child) => child.update_child_difference(),
};
self.update_child_difference();
}
pub fn subchild_to_parent(&mut self, index: (usize, usize)) {
(
self.items,
self.children[index.0].as_deref_mut().unwrap().children[index.1],
self.children[index.0].as_deref_mut().unwrap().children_size[index.1],
self.children[index.1].as_deref_mut().unwrap().children[index.0],
self.children[index.1].as_deref_mut().unwrap().children_size[index.0],
) = match self.children[index.1].as_deref_mut().unwrap().children[index.0].as_deref_mut() {
None => Default::default(),
Some(child) => (
child.take_items(),
child.take_child(index.0),
child.take_child_size(index.0),
child.take_child(index.1),
child.take_child_size(index.1),
),
};
self.children[index.1]
.as_deref_mut()
.unwrap()
.update_child_difference();
self.update_child_difference();
}
pub fn update_child_difference(&mut self) {
let left = match &self.children[0] {
Some(child) => {
self.children_size[0] = child.items.iter().filter(|item| item.is_some()).count()
+ child.children_size[0]
+ child.children_size[1];
if child.is_full() {
child.child_difference + 1
} else {
0
}
}
None => 0,
};
let right = match &self.children[1] {
Some(child) => {
self.children_size[1] = child.items.iter().filter(|item| item.is_some()).count()
+ child.children_size[0]
+ child.children_size[1];
if child.is_full() {
child.child_difference + 1
} else {
0
}
}
None => 0,
};
self.child_difference = left.abs_diff(right);
}
pub fn height_balance(&mut self) {
if !self.is_height_balanced() {
match self.update_height_difference().1 {
Less => self.double_rotation(Side::Left),
Equal => {}
Greater => self.double_rotation(Side::Right),
}
}
}
pub fn is_height_balanced(&mut self) -> bool {
self.update_height();
self.update_height_difference().0 < 2
}
pub fn update_height(&mut self) {
let left = match &self.children[0] {
None => 0,
Some(child) => child.height + 1,
};
let right = match &self.children[1] {
None => 0,
Some(child) => child.height + 1,
};
self.height = max(left, right);
}
pub fn update_height_difference(&mut self) -> (u8, Ordering) {
self.update_height();
let left_height: u8 = match &self.children[0] {
None => 0,
Some(child) => child.height,
};
let right_height: u8 = match &self.children[1] {
None => 0,
Some(child) => child.height,
};
(
left_height.abs_diff(right_height),
left_height.cmp(&right_height),
)
}
pub fn double_rotation(&mut self, side: Side) {
let index = match side {
Side::Left => (0, 1),
Side::Right => (1, 0),
};
let new_child = {
let mut node = Node::new();
node.items = self.take_items();
node.children[index.0] = self.take_child(index.0);
node.children_size[index.0] = self.take_child_size(index.0);
(node.children[index.1], node.children_size[index.1]) = {
let node = self.children[index.1].as_deref_mut().unwrap();
(node.take_child(index.0), node.take_child_size(index.0))
};
node.update_child_difference();
node.update_height_difference();
node
};
let mut node = self.take_child(index.1).unwrap();
self.items = node.take_items();
self.children[index.1] = node.take_child(index.1);
self.children_size[index.1] = node.take_child_size(index.1);
self.children[index.0] = Some(new_child);
self.children_size[index.0] = match self.children[index.0].as_deref() {
Some(node) => node.children_size[0] + node.children_size[1] + ARRAY_SIZE,
None => 0,
};
self.update_child_difference();
node.update_height_difference();
}
}