Binary Search Tree

Rust Implementation of Binary Search Tree


//! Author: https://github.com/dhruvildave

/// Binary Search Tree
pub(crate) mod binary_search_tree {
    use std::cell::RefCell;
    use std::fmt;
    use std::fmt::Debug;
    use std::rc::Rc;

    #[derive(Clone, Debug)]
    pub(crate) struct InsertDuplicateError<T>(T);

    impl<T> fmt::Display for InsertDuplicateError<T>
    where
        T: Clone + Debug,
    {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            write!(f, "trying to insert duplicate value: {:?}", self.0)
        }
    }

    pub(crate) type InsertResult<T> = Result<(), InsertDuplicateError<T>>;

    #[derive(Debug)]
    pub(crate) struct TreeNode<T>
    where
        T: PartialOrd,
    {
        val: T,
        left: Tree<T>,
        right: Tree<T>,
    }

    #[derive(Debug)]
    pub(crate) struct Tree<T>
    where
        T: PartialOrd,
    {
        root: Option<Rc<RefCell<TreeNode<T>>>>,
    }

    impl<T> Tree<T>
    where
        T: PartialOrd,
    {
        pub(crate) fn new(val: T) -> Self {
            Self {
                root: Some(Rc::new(RefCell::new(TreeNode::new(val)))),
            }
        }

        /// Recursively searches for the value.
        pub(crate) fn search(&self, val: &T) -> Option<Rc<RefCell<TreeNode<T>>>> {
            let mut c = self.root.as_ref().cloned();

            while let Some(n) = c {
                if n.borrow().val == *val {
                    return Some(n);
                }

                if *val < n.borrow().val {
                    c = n.borrow().left.root.as_ref().cloned();
                } else {
                    c = n.borrow().right.root.as_ref().cloned();
                }
            }

            None
        }

        /// Inserts a value in the BST. If the value exists, it throws an errors.
        pub(crate) fn insert(&mut self, val: T) -> InsertResult<T> {
            let mut node = self.root.as_mut().unwrap().borrow_mut();
            if val < node.val {
                // Insert into the left subtree
                match node.left.root {
                    Some(_) => node.left.insert(val),
                    None => {
                        node.left = Tree::new(val);
                        Ok(())
                    }
                }
            } else if val > node.val {
                // Insert into the right subtree
                match node.right.root {
                    Some(_) => node.right.insert(val),
                    None => {
                        node.right = Tree::new(val);
                        Ok(())
                    }
                }
            } else {
                Err(InsertDuplicateError(val))
            }
        }

        /// Calculate height of the BST.
        pub(crate) fn height(&self) -> usize {
            if let Some(ref root) = self.root {
                1 + match (&root.borrow().left.root, &root.borrow().right.root) {
                    (None, None) => 0,
                    (Some(_), None) => root.borrow().left.height(),
                    (None, Some(_)) => root.borrow().right.height(),
                    (Some(_), Some(_)) => {
                        usize::max(root.borrow().left.height(), root.borrow().right.height())
                    }
                }
            } else {
                0
            }
        }

        /// [Self::inorder_walk] helper function.
        fn inorder_rec(&self, vec: &mut Vec<T>)
        where
            T: Clone,
        {
            if let Some(ref r) = self.root {
                r.borrow().left.inorder_rec(vec);
                vec.push(r.borrow().value());
                r.borrow().right.inorder_rec(vec);
            }
        }

        /// Inorder Traversal.
        pub(crate) fn inorder_walk(&self) -> Vec<T>
        where
            T: Clone,
        {
            let mut vec = Vec::new();
            self.inorder_rec(&mut vec);
            vec
        }

        /// [Self::preorder_walk] helper function.
        fn preorder_rec(&self, vec: &mut Vec<T>)
        where
            T: Clone,
        {
            if let Some(ref r) = self.root {
                vec.push(r.borrow().value());
                r.borrow().left.preorder_rec(vec);
                r.borrow().right.preorder_rec(vec);
            }
        }

        /// Preorder Traversal.
        pub(crate) fn preorder_walk(&self) -> Vec<T>
        where
            T: Clone,
        {
            let mut vec = Vec::new();
            self.preorder_rec(&mut vec);
            vec
        }

        /// [Self::postorder_walk] helper function.
        fn postorder_rec(&self, vec: &mut Vec<T>)
        where
            T: Clone,
        {
            if let Some(ref r) = self.root {
                r.borrow().left.postorder_rec(vec);
                r.borrow().right.postorder_rec(vec);
                vec.push(r.borrow().value());
            }
        }

        /// Postorder Traversal.
        pub(crate) fn postorder_walk(&self) -> Vec<T>
        where
            T: Clone,
        {
            let mut vec = Vec::new();
            self.postorder_rec(&mut vec);
            vec
        }

        /// Returns an option to the minimum value in the tree.
        pub(crate) fn min(&self) -> Option<T>
        where
            T: Clone,
        {
            self.root.as_ref()?;

            let node = self.root.as_ref().unwrap().borrow();
            match node.left.root {
                Some(_) => node.left.min(),
                None => Some(node.value()),
            }
        }

        /// Returns an option to the maximum value in the tree.
        pub(crate) fn max(&self) -> Option<T>
        where
            T: Clone,
        {
            self.root.as_ref()?;

            let node = self.root.as_ref().unwrap().borrow();
            match node.right.root {
                Some(_) => node.right.max(),
                None => Some(node.value()),
            }
        }

        /// [Self::delete] helper function.
        fn delete_rec(&self, val: &T) -> Option<Rc<RefCell<TreeNode<T>>>>
        where
            T: Clone,
        {
            if let Some(ref n) = self.root {
                if n.borrow().val > *val {
                    let l = n.borrow().left.delete_rec(val);
                    n.borrow_mut().left.root = l;
                } else if n.borrow().val < *val {
                    let r = n.borrow().right.delete_rec(val);
                    n.borrow_mut().right.root = r;
                } else {
                    if n.borrow().left.root.is_none() {
                        return n.borrow().right.root.as_ref().cloned();
                    }
                    if n.borrow().right.root.is_none() {
                        return n.borrow().left.root.as_ref().cloned();
                    }

                    // Recursively replacing the minimum value of the right sub-tree.
                    let next = n.borrow().right.min();
                    if let Some(value) = next {
                        let r = n.borrow().right.delete_rec(&value);
                        n.borrow_mut().val = value;
                        n.borrow_mut().right.root = r;
                    }
                }
            }

            self.root.as_ref().cloned()
        }

        /// Delete a value from the BST.
        pub(crate) fn delete(&mut self, val: &T)
        where
            T: Clone,
        {
            self.root = self.delete_rec(val);
        }

        /// [Self::is_balanced] helper function.
        fn is_balanced_rec(&self) -> (bool, usize) {
            if self.root.is_none() {
                return (true, 0);
            }

            let root = self.root.as_ref().unwrap();

            let left = root.borrow().left.is_balanced_rec();
            let right = root.borrow().right.is_balanced_rec();

            let balanced = left.0 && right.0 && usize::abs_diff(left.1, right.1) <= 1;
            (balanced, 1 + usize::max(left.1, right.1))
        }

        /// Check if BST is balanced or not.
        pub(crate) fn is_balanced(&self) -> bool {
            self.is_balanced_rec().0
        }
    }

    impl<T> TreeNode<T>
    where
        T: PartialOrd,
    {
        fn new(val: T) -> Self {
            Self {
                val,
                left: Tree { root: None },
                right: Tree { root: None },
            }
        }

        fn value(&self) -> T
        where
            T: Clone,
        {
            self.val.clone()
        }
    }
}

fn main() {
    use crate::binary_search_tree::Tree;

    let mut t = Tree::new(&10);
    for i in [&5, &15, &13, &16, &2, &8] {
        t.insert(i).unwrap();
    }

    dbg!(t.height());
    dbg!(t.min());
    dbg!(t.max());
    dbg!(t.search(&&12));

    t.delete(&&10);

    dbg!(t.preorder_walk());
    dbg!(t.inorder_walk());
    dbg!(t.postorder_walk());
    dbg!(t.is_balanced());
}