LeetCode 669

leetcode669.rs

// This is a solution to a modified version of LeetCode 669. The original problem
// requires the modified root of the binary tree to be returned, while the solution
// below modifies it in place rather than returning the new root.
//
// https://leetcode.com/problems/trim-a-binary-search-tree/
// https://rust.cc/article?id=93bb5260-7e17-4d1c-bc42-08a9873f219c

use std::cell::RefCell;
use std::rc::Rc;

impl Solution {
    pub fn trim_bst(root: &mut Option<Rc<RefCell<TreeNode>>>, l: i32, r: i32) {
        if let Some(inside) = root.take() {
            let mut n = inside.borrow_mut();
            if n.val < l {
                Self::trim_bst(&mut n.right, l, r);
            } else if n.val > r {
                Self::trim_bst(&mut n.left, l, r);
            } else {
                Self::trim_bst(&mut n.left, l, r);
                Self::trim_bst(&mut n.right, l, r);
                drop(n);
                *root = Some(inside);
            }
        }
    }
}

fn main() {
    let mut t = TreeNode::full_from_arr(vec![Some(1), Some(0), Some(2)], 0);
    Solution::trim_bst(&mut t, 1, 2);
    println!("{:#?}", t);
}

// deps
pub struct Solution;
#[derive(Debug, PartialEq, Eq)]
pub struct TreeNode {
    pub val: i32,
    pub left: Option<Rc<RefCell<TreeNode>>>,
    pub right: Option<Rc<RefCell<TreeNode>>>,
}

impl TreeNode {
    #[inline]
    pub fn new(val: i32) -> Self {
        TreeNode {
            val,
            left: None,
            right: None,
        }
    }

    pub fn full_from_arr(arr: Vec<Option<i32>>, i: usize) -> Option<Rc<RefCell<TreeNode>>> {
        if arr.len() == 0 || i >= arr.len() || arr[i] == None {
            return None;
        }

        let root = Rc::new(RefCell::new(TreeNode::new(arr[i].unwrap())));
        root.borrow_mut().left = Self::full_from_arr(arr.clone(), 2 * i + 1);
        root.borrow_mut().right = Self::full_from_arr(arr.clone(), 2 * i + 2);

        Some(root)
    }
}