Rust-AVL木

avl_tree.rs

use std::boxed::Box;
use std::cmp::max;

#[derive(Debug, Clone, PartialEq)]
pub struct Node {
    height: isize,
    key: isize,
    left: Option<Box<Node>>,
    right: Option<Box<Node>>,
}


#[derive(Debug, Clone, PartialEq)]
pub struct AVL {
    root: Option<Box<Node>>,
}


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;
    }
}