use std::boxed::Box; use std::cmp::max; #[derive(Debug, Clone, PartialEq)] pub struct Node { height: isize, key: isize, left: Option>, right: Option>, } #[derive(Debug, Clone, PartialEq)] pub struct AVL { root: Option>, } impl AVL { pub fn new() -> Self { Self { root: None } } pub fn insert(&mut self, insert_key: isize) { if let Some(node) = &mut self.root { node.insert(insert_key); node.modify_height(); } else { // 空のとき入れる let key = Box::new(Node::new(insert_key)); self.root = Some(key); } } pub fn search(&self, search_key: isize) -> bool { if let Some(node) = &self.root { node.search(search_key) } else { false } } pub fn delete(&mut self, delete_key: isize) -> bool { if let Some(node) = &mut self.root { node.delete(delete_key) } else { false } } } impl Node { fn new(n_key: isize) -> Self { Self { height: 1, key: n_key, left: None, right: None, } } fn insert(&mut self, insert_key: isize) { if self.key < insert_key { if let Some(node) = &mut self.right { node.insert(insert_key); node.modify_height(); } else { let key = Box::new(Self::new(insert_key)); self.right = Some(key); self.modify_height(); } if self.check_height() { self.rotate_r(); } } else if self.key > insert_key { // self.key > insert_key if let Some(node) = &mut self.left { node.insert(insert_key); node.modify_height(); } else { let key = Box::new(Self::new(insert_key)); self.left = Some(key); self.modify_height(); } if self.check_height() { self.rotate_l(); } } else { // Do nothing } self.modify_height(); } fn search(&self, search_key: isize) -> bool { if self.key < search_key { if let Some(node) = &self.right { node.search(search_key) } else { false } } else if self.key > search_key { // self.key > insert_key if let Some(node) = &self.left { node.search(search_key) } else { false } } else { true } } fn delete(&mut self, delete_key: isize) -> bool // "Hit"=>true, "Nothing"=>false { if self.key < delete_key { if let Some(node) = &mut self.right { if node.key == delete_key { // Hit! if let Some(node_left) = &mut node.left { let max_key = node_left.delete_max_key(); node.delete(max_key); node.key = max_key; if node.check_height() { node.rotate_r(); } } else if let Some(node_right) = &mut node.right { let key = node_right.key; node.delete(key); node.key = key; } else { self.right = None; } self.modify_height(); true } else { node.delete(delete_key) } } else { false } } else if self.key > delete_key { // self.key > insert_key if let Some(node) = &mut self.left { if node.key == delete_key { // Hit! if let Some(node_left) = &mut node.left { let max_key = node_left.delete_max_key(); node.delete(max_key); node.key = max_key; if node.check_height() { node.rotate_l(); } } else if let Some(node_right) = &mut node.right { let key = node_right.key; node.delete(key); node.key = key; } else { self.left = None; } self.modify_height(); true } else { node.delete(delete_key) } } else { false } } else { false } } } impl Node { fn delete_max_key(&mut self) -> isize // { if let Some(node_right) = &mut self.right { node_right.delete_max_key() } else if let Some(node_left) = &mut self.left { node_left.delete_max_key() } else { self.key } } fn modify_height(&mut self) { let (left, right) = { let r = if let Some(n) = &self.right { n.height } else { 0 }; let l = if let Some(n) = &self.left { n.height } else { 0 }; (l, r) }; self.height = max(left, right) + 1; } fn check_height(&self) -> bool { let (left, right) = { let r = if let Some(n) = &self.right { n.height } else { 0 }; let l = if let Some(n) = &self.left { n.height } else { 0 }; (l, r) }; let diff = (left - right).abs(); diff == 2 } fn rotate_l(&mut self) { let mut shaft_node_key = self.clone(); let mut l = self.left.clone().unwrap(); let (left, right) = { let rh = if let Some(n) = &l.right { n.height } else { 0 }; let lh = if let Some(n) = &l.left { n.height } else { 0 }; (lh, rh) }; if left > right { // l.left > l.right if let Some(l_right) = &mut l.right { shaft_node_key.left = Some(l_right.clone()); shaft_node_key.modify_height(); *l_right = Box::new(shaft_node_key); } else { shaft_node_key.left = None; shaft_node_key.modify_height(); l.right = Some(Box::new(shaft_node_key)); } } else { // l.left < l.right if let Some(l_right) = &mut l.right { shaft_node_key.left = None; l_right.left = Some(Box::new({ let mut n = Self::new(l.key); n.modify_height(); n })); l_right.right = Some(Box::new({ shaft_node_key.modify_height(); shaft_node_key })); l_right.modify_height(); l = l_right.clone(); } } *self = *l; } fn rotate_r(&mut self) { let mut shaft_node_key = self.clone(); let mut r = self.right.clone().unwrap(); let (left, right) = { let rh = if let Some(n) = &r.right { n.height } else { 0 }; let lh = if let Some(n) = &r.left { n.height } else { 0 }; (lh, rh) }; if left < right { // r.left < r.right if let Some(r_left) = &mut r.left { shaft_node_key.right = Some(r_left.clone()); shaft_node_key.modify_height(); *r_left = Box::new(shaft_node_key); } else { shaft_node_key.right = None; shaft_node_key.modify_height(); r.left = Some(Box::new(shaft_node_key)); } } else { // r.left > r.right if let Some(r_left) = &mut r.left { shaft_node_key.left = None; r_left.right = Some(Box::new({ let mut n = Self::new(r.key); n.modify_height(); n })); r_left.left = Some(Box::new({ shaft_node_key.modify_height(); shaft_node_key })); r_left.modify_height(); r = r_left.clone(); } } *self = *r; } }