3
\$\begingroup\$

The original problem statement is here, here is the super short summary:

Your binary search tree iterator must implement the following interface:

pub struct TreeNode {
  pub val: i32,
  pub left: Option<Rc<RefCell<TreeNode>>>,
  pub right: Option<Rc<RefCell<TreeNode>>>,
}

struct BSTIterator {}

impl BSTIterator {
    fn new(root: Option<Rc<RefCell<TreeNode>>>) -> Self {}
    
    fn next(&mut self) -> i32 {}
    
    fn has_next(&self) -> bool {}
}

The binary search iterator must return the values of the binary search tree from smallest to largest. It must use O(h) memory, where h is the height of the tree. Both next and has_next must run in O(1) amortized time.

My solution to this is gradual, controlled recursion:

  1. Given the root of the tree in constructor, traverse all the left children and remember them.
  2. On each next call, pop the last remembered node and return it. If that node had a right child, traverse it and all of its left children, and remember them.

The implementation:

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

// Definition for a binary tree node.
#[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,
        }
    }
}

struct BSTIterator {
    nodes: Vec<Rc<RefCell<TreeNode>>>,
}

impl BSTIterator {
    fn new(root: Option<Rc<RefCell<TreeNode>>>) -> Self {
        let mut nodes = vec![];

        if let Some(root) = root {
            BSTIterator::insert(&mut nodes, root);
        }

        BSTIterator { nodes }
    }

    /// Return the next smallest number
    fn next(&mut self) -> i32 {
        let rc = self.nodes.pop().unwrap();
        let node = (*rc).borrow();

        if let Some(node) = &node.right {
            BSTIterator::insert(&mut self.nodes, Rc::clone(node));
        }

        node.val
    }

    /// True if there is a next smallest number, false otherwise
    fn has_next(&self) -> bool {
        !self.nodes.is_empty()
    }

    fn insert(nodes: &mut Vec<Rc<RefCell<TreeNode>>>, node: Rc<RefCell<TreeNode>>) {
        let mut node = Some(node);
        while let Some(rc) = node {
            nodes.push(Rc::clone(&rc));

            node = (*rc)
                .borrow()
                .left
                .as_ref()
                .and_then(|lft| Some(Rc::clone(&lft)));
        }
    }
}

fn main() {
    let mut node50 = TreeNode::new(50);
    let mut node25 = TreeNode::new(25);
    let mut node75 = TreeNode::new(75);
    let mut node12 = TreeNode::new(12);
    let mut node32 = TreeNode::new(32);
    let mut node62 = TreeNode::new(62);
    let mut node80 = TreeNode::new(80);

    node25.left = Some(Rc::new(RefCell::new(node12)));
    node25.right = Some(Rc::new(RefCell::new(node32)));

    node75.left = Some(Rc::new(RefCell::new(node62)));
    node75.right = Some(Rc::new(RefCell::new(node80)));

    node50.left = Some(Rc::new(RefCell::new(node25)));
    node50.right = Some(Rc::new(RefCell::new(node75)));

    let mut it = BSTIterator::new(Some(Rc::new(RefCell::new(node50))));

    while it.has_next() {
        println!("{}", it.next());
    }
}

The compiler kept screaming at me, and it was my first time working with Rc and RefCell. Which anti-patterns do I have here? Is (*rc).borrow().left.as_ref().and_then(/*...*/) a reasonable thing to write in insert?

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.