Binary Search Tree - Rust

main.rs

#![allow(unused_variables, dead_code)]

use std::fmt::{Display, Formatter};

#[derive(Clone, PartialEq, Eq)]
struct BST {
    key: String,
    value: u64,
    left: Box<Option<BST>>,
    right: Box<Option<BST>>
}

impl BST {
    fn height(&self) -> usize {
        if self.left.is_some() & self.right.is_some() {
            let left_height = &self.left.clone().unwrap().height();
            let right_height = &self.right.clone().unwrap().height();
            return 1 + std::cmp::max(left_height, right_height);
        } else if self.left.is_some() {
            let left_height = &self.left.clone().unwrap().height();
            return 1 + left_height;
        } else if self.right.is_some() {
            let right_height = &self.right.clone().unwrap().height();
            return 1 + right_height;
        } else {
            return 1;
        }
    }

    fn node_count(&self) -> usize {
        let mut bst_count: usize = 1;
        if self.left.is_some() {
            bst_count += &self.left.clone().unwrap().node_count();
        }
        if self.right.is_some() {
            bst_count += &self.right.clone().unwrap().node_count();
        }
        bst_count
    }

    fn nodes_on_level(&self, current_level: usize, desired_level: usize) -> Vec<Option<BST>> {
        let mut vec: Vec<Option<BST>> = Vec::new();
        if current_level == desired_level {
            vec.push(Some(self.clone()));
        } else {
            if self.left.is_some() {
                let left_vec: Vec<Option<BST>> = self.left.clone()
                                            .unwrap()
                                            .nodes_on_level(current_level + 1, desired_level);
                vec.extend_from_slice(&left_vec);
            } else {
                vec.push(None)
            }
            if self.right.is_some() {
                let right_vec: Vec<Option<BST>> = self.right.clone()
                                            .unwrap()
                                            .nodes_on_level(current_level + 1, desired_level);
                vec.extend_from_slice(&right_vec);
            } else {
                vec.push(None)
            }
        }
        vec
    }

    fn max_iterative(&self) -> BST {
        if self.right.is_some() {
            let mut right = &self.right.clone().unwrap();
            loop {
                match right.right.as_ref() {
                    Some(bst) => {
                        right = bst;
                    },
                    None => break
                }
            }
            return right.clone();
        }
        self.clone()
    }

    fn max_recursive(&self) -> BST {
        if self.right.is_some() {
            return self.right.clone().unwrap().max_recursive();
        }
        self.clone()
    }

    fn min_iterative(&self) -> BST {
        if self.left.is_some() {
            let mut left = &self.left.clone().unwrap();
            loop {
                match left.left.as_ref() {
                    Some(bst) => {
                        left = bst;
                    },
                    None => break
                }
            }
            return left.clone();
        }
        self.clone()
    }

    fn min_recursive(&self) -> BST {
        if self.left.is_some() {
            return self.left.clone().unwrap().min_recursive();
        }
        self.clone()
    }

    fn insert(&mut self, bst: BST) -> BST {
        if self.value < bst.value {
            if self.right.is_some() {
                self.right = Box::new(Some(self.right.clone().unwrap().insert(bst)));
            } else {
                self.right = Box::new(Some(bst))
            }
        } else if self.value > bst.value {
            if self.left.is_some() {
                self.left = Box::new(Some(self.left.clone().unwrap().insert(bst)));
            } else {
                self.left = Box::new(Some(bst))
            }
        }
        self.clone()
    }

    fn remove(&mut self, value: u64) -> Option<BST> {
        if self.value == value {
            return None;
        } else if self.value < value {
            if self.right.is_some() {
                let mut right = self.right.clone().unwrap();
                if right.value == value {
                    self.right = Box::new(None);
                } else {
                    self.right = Box::new(right.remove(value));
                }
            }
        } else if self.value > value {
            if self.left.is_some() {
                let mut left = self.left.clone().unwrap();
                if left.value == value {
                    self.left = Box::new(None);
                } else {
                    self.left = Box::new(left.remove(value));
                }
            }
        }
        Some(self.clone())
    }

    fn balance(&mut self) {
        // No algorithms included, simply convert BST to vector and convert the result back to BST again.
        let values: Vec<(String, u64)> = self.clone().into();
        *self = BST::from(values);
    }

    fn recursive_search(&self, value: u64) -> Option<BST> {
        if self.value == value {
            return Some(self.clone());
        } else if (self.value < value) & self.right.is_some() {
            return self.right.clone().unwrap().recursive_search(value);
        } else if (self.value > value) & self.left.is_some() {
            return self.left.clone().unwrap().recursive_search(value);
        } else {
            return None;
        }
    }

    fn iterative_search(&self, value: u64) -> Option<BST> {
        let mut bst = self.clone();
        loop {
            if (value > bst.value) & bst.right.is_some() {
                bst = self.right.clone().unwrap();
            } else if (value < bst.value) & bst.left.is_some() {
                bst = self.left.clone().unwrap();
            } else if value == bst.value {
                return Some(bst);
            } else {
                return None;
            }
        }
    }

    fn inorder_tree_walk(&self) {
        if self.left.is_some() {
            self.left.clone().unwrap().inorder_tree_walk();
        }
        println!("{}", self);
        if self.right.is_some() {
            self.right.clone().unwrap().inorder_tree_walk();
        }
    }

    fn preorder_tree_walk(&self) {
        println!("{}", self);
        if self.left.is_some() {
            self.left.clone().unwrap().preorder_tree_walk();
        }
        if self.right.is_some() {
            self.right.clone().unwrap().preorder_tree_walk();
        }
    }

    fn postorder_tree_walk(&self) {
        if self.left.is_some() {
            self.left.clone().unwrap().postorder_tree_walk();
        }
        if self.right.is_some() {
            self.right.clone().unwrap().postorder_tree_walk();
        }
        println!("{}", self);
    }

    fn print(&self) {
        let height = self.height();
        for level in 0..height {
            let node_list = self.nodes_on_level(0, level);
            let floor = height - level;
            let vertices_count = u64::pow(2, floor as u32) as usize;
            let space_left = u64::pow(2, floor as u32) as usize - 1;
            let space_between = u64::pow(2, floor as u32) as usize - 1;

            let mut print_node_str = " ".repeat(2 * space_left);
            for node in &node_list {
                match node {
                    Some(n) => {
                        print_node_str.push_str(&format!("{}", n));
                    },
                    None => {
                        print_node_str.push_str(&" ".repeat(5));
                    }
                }
                print_node_str.push_str(&" ".repeat(3 * space_between));
            }
            println!("{}", print_node_str);

            for i in 1..vertices_count {
                let mut print_vertices_str = " ".repeat(2 * space_left - i + 3).to_string();
                for (j, node) in node_list.iter().enumerate() {
                    if node.is_none() {
                        print_vertices_str.push_str(&" ".repeat(2 * vertices_count - i));
                        continue
                    }
                    let _node = node.as_ref().unwrap();
                    let left = _node.left.is_some();
                    let right = _node.right.is_some();
                    if left {
                        print_vertices_str.push_str("/");
                    } else {
                        print_vertices_str.push_str(" ");
                    }
                    print_vertices_str.push_str(&" ".repeat(2 * i + 2 * j));
                    if right {
                        print_vertices_str.push_str("\\");
                    } else {
                        print_vertices_str.push_str(" ");
                    }
                    print_vertices_str.push_str(&" ".repeat(3 * vertices_count - 2 * i));
                }
                println!("{}", print_vertices_str);
            }
        }
    }
}

impl Display for BST {
    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
        write!(f, "({:?}:{})", self.key.clone(), self.value)
    }
}

impl From<Vec<(String, u64)>> for BST {
    fn from(values: Vec<(String, u64)>) -> Self {
        if values.len() == 0 {
            return BST {
                key: "".to_string(),
                value: 0,
                left: Box::new(None),
                right: Box::new(None)
            }
        }
        let mut sorted = values;
        sorted.sort_by_key(|k| k.1);
        let middle_index = sorted.len() / 2;
        let mut right_vec = sorted.split_off(middle_index);
        sorted.reverse();
        let middle_entry = right_vec.remove(0);
        let mut bst = BST {
            key: middle_entry.0,
            value: middle_entry.1,
            left: Box::new(None),
            right: Box::new(None)
        };
        for entry in sorted {
            bst.insert(BST {
                key: entry.0,
                value: entry.1,
                left: Box::new(None),
                right: Box::new(None)
            });
        }
        for entry in right_vec {
            bst.insert(BST {
                key: entry.0,
                value: entry.1,
                left: Box::new(None),
                right: Box::new(None)
            });
        }
        bst
    }
}

impl Into<Vec<(String, u64)>> for BST {
    fn into(self) -> Vec<(String, u64)> {
        let mut values: Vec<(String, u64)> = Vec::new();
        values.push((self.key.clone(), self.value));
        if self.left.is_some() {
            let left_vec: Vec<(String, u64)> = self.left.clone().unwrap().into();
            values.extend_from_slice(&left_vec);
        }
        if self.right.is_some() {
            let right_vec: Vec<(String, u64)> = self.right.clone().unwrap().into();
            values.extend_from_slice(&right_vec);
        }
        values
    }
}

fn main() {
    let test = BST::from(vec![("T".to_string(), 0), ("e".to_string(), 1),
                            ("s".to_string(), 2), ("t".to_string(), 3), ("!".to_string(), 4)]);
    println!("Tree view:");
    test.print();

    println!("Inorder Tree Walk:");
    test.inorder_tree_walk();

    println!("Preorder Tree Walk:");
    test.preorder_tree_walk();

    println!("Postorder Tree Walk:");
    test.postorder_tree_walk();
}


#[cfg(test)]
mod bst_tests {
    use super::*;

    #[test]
    fn test_bst_from_vec() {
        let values = vec![
            ("A".to_string(), 10),
            ("C".to_string(), 11),
            ("B".to_string(), 9),
            ("E".to_string(), 8),
            ("D".to_string(), 13)
        ];
        let bst = BST::from(values);
        assert!(bst == BST {
            key: "A".to_string(),
            value: 10,
            left: Box::new(Some(BST {
                key: "B".to_string(),
                value: 9,
                left: Box::new(Some(BST {
                    key: "E".to_string(),
                    value: 8,
                    left: Box::new(None),
                    right: Box::new(None)
                })),
                right: Box::new(None)
            })),
            right: Box::new(Some(BST {
                key: "C".to_string(),
                value: 11,
                left: Box::new(None),
                right: Box::new(Some(BST {
                    key: "D".to_string(),
                    value: 13,
                    left: Box::new(None),
                    right: Box::new(None)
                }))
            }))
        })
    }

    #[test]
    fn test_bst_into_vec() {
        let bst = BST {
            key: "A".to_string(),
            value: 10,
            left: Box::new(Some(BST {
                key: "B".to_string(),
                value: 9,
                left: Box::new(None),
                right: Box::new(None)
            })),
            right: Box::new(Some(BST {
                key: "C".to_string(),
                value: 11,
                left: Box::new(None),
                right: Box::new(None)
            }))
        };
        let values: Vec<(String, u64)> = bst.into();
        assert!(values == vec![("A".to_string(), 10), ("B".to_string(), 9), ("C".to_string(), 11)]);
    }

    #[test]
    fn test_bst_nodes_on_level() {
        let values = vec![
            ("A".to_string(), 10),
            ("C".to_string(), 11),
            ("B".to_string(), 9)
        ];
        let bst = BST::from(values);
        assert!(bst.nodes_on_level(0, 1) == vec![Some(BST {
            key: "B".to_string(),
            value: 9,
            left: Box::new(None),
            right: Box::new(None)
        }), Some(BST {
            key: "C".to_string(),
            value: 11,
            left: Box::new(None),
            right: Box::new(None)
        })])
    }

    #[test]
    fn test_bst_insert() {
        let mut bst = BST {
            key: "A".to_string(),
            value: 10,
            left: Box::new(None),
            right: Box::new(None)
        };
        bst.insert(BST {
            key: "C".to_string(),
            value: 11,
            left: Box::new(None),
            right: Box::new(None)
        });
        bst.insert(BST {
            key: "B".to_string(),
            value: 9,
            left: Box::new(None),
            right: Box::new(None)
        });
        assert!(bst == BST {
            key: "A".to_string(),
            value: 10,
            left: Box::new(Some(BST {
                key: "B".to_string(),
                value: 9,
                left: Box::new(None),
                right: Box::new(None)
            })),
            right: Box::new(Some(BST {
                key: "C".to_string(),
                value: 11,
                left: Box::new(None),
                right: Box::new(None)
            }))
        })
    }

    #[test]
    fn test_bst_remove() {
        let mut bst = BST {
            key: "A".to_string(),
            value: 10,
            left: Box::new(Some(BST {
                key: "B".to_string(),
                value: 9,
                left: Box::new(None),
                right: Box::new(None)
            })),
            right: Box::new(Some(BST {
                key: "C".to_string(),
                value: 11,
                left: Box::new(None),
                right: Box::new(None)
            }))
        };
        bst.remove(9);
        assert!(bst == BST {
            key: "A".to_string(),
            value: 10,
            left: Box::new(None),
            right: Box::new(Some(BST {
                key: "C".to_string(),
                value: 11,
                left: Box::new(None),
                right: Box::new(None)
            }))
        })
    }

    #[test]
    fn test_bst_height() {
        let mut bst = BST {
            key: "A".to_string(),
            value: 10,
            left: Box::new(None),
            right: Box::new(None)
        };
        assert!(bst.height() == 1);
        bst.insert(BST {
            key: "C".to_string(),
            value: 11,
            left: Box::new(None),
            right: Box::new(None)
        });
        assert!(bst.height() == 2);
        bst.insert(BST {
            key: "B".to_string(),
            value: 9,
            left: Box::new(None),
            right: Box::new(None)
        });
        assert!(bst.height() == 2);
        bst.insert(BST {
            key: "D".to_string(),
            value: 8,
            left: Box::new(Some(BST {
                key: "E".to_string(),
                value: 7,
                left: Box::new(None),
                right: Box::new(None)
            })),
            right: Box::new(None)
        });
        assert!(bst.height() == 4);
        bst.remove(7);
        assert!(bst.height() == 3);
    }

    #[test]
    fn test_bst_node_count() {
        let mut bst = BST {
            key: "A".to_string(),
            value: 10,
            left: Box::new(None),
            right: Box::new(None)
        };
        assert!(bst.node_count() == 1);
        bst.insert(BST {
            key: "C".to_string(),
            value: 11,
            left: Box::new(None),
            right: Box::new(None)
        });
        assert!(bst.node_count() == 2);
        bst.insert(BST {
            key: "B".to_string(),
            value: 9,
            left: Box::new(None),
            right: Box::new(None)
        });
        assert!(bst.node_count() == 3);
        bst.insert(BST {
            key: "D".to_string(),
            value: 8,
            left: Box::new(Some(BST {
                key: "E".to_string(),
                value: 7,
                left: Box::new(None),
                right: Box::new(None)
            })),
            right: Box::new(None)
        });
        assert!(bst.node_count() == 5);
        bst.remove(7);
        assert!(bst.node_count() == 4);
    }

    #[test]
    fn test_bst_balance() {
        let mut bst = BST {
            key: "A".to_string(),
            value: 10,
            left: Box::new(Some(BST {
                key: "B".to_string(),
                value: 9,
                left: Box::new(None),
                right: Box::new(Some(BST {
                    key: "C".to_string(),
                    value: 11,
                    left: Box::new(None),
                    right: Box::new(None)
                }))
            })),
            right: Box::new(None)
        };
        bst.balance();
        assert!(bst == BST {
            key: "A".to_string(),
            value: 10,
            left: Box::new(Some(BST {
                key: "B".to_string(),
                value: 9,
                left: Box::new(None),
                right: Box::new(None)
            })),
            right: Box::new(Some(BST {
                key: "C".to_string(),
                value: 11,
                left: Box::new(None),
                right: Box::new(None)
            }))
        })
    }
}