919. Complete Binary Tree Inserter

A complete binary tree is a binary tree in which every level, except possibly the last, is completely filled, and all nodes are as far left as possible.

Write a data structure CBTInserter that is initialized with a complete binary tree and supports the following operations:

  • CBTInserter(TreeNode root) initializes the data structure on a given tree with head node root;
  • CBTInserter.insert(int v) will insert a TreeNode into the tree with value node.val = v so that the tree remains complete, and returns the value of the parent of the inserted TreeNode;
  • CBTInserter.get_root() will return the head node of the tree.

 

Example 1:

Input: inputs = ["CBTInserter","insert","get_root"], inputs = [[[1]],[2],[]]
Output: [null,1,[1,2]]

Example 2:

Input: inputs = ["CBTInserter","insert","insert","get_root"], inputs = [[[1,2,3,4,5,6]],[7],[8],[]]
Output: [null,3,4,[1,2,3,4,5,6,7,8]]

 

Note:

  1. The initial given tree is complete and contains between 1 and 1000 nodes.
  2. CBTInserter.insert is called at most 10000 times per test case.
  3. Every value of a given or inserted node is between 0 and 5000.

 

 

Rust Solution

use rustgym_util::*;
struct CBTInserter {
    stack: Vec<TreeLink>,
}

impl CBTInserter {
    fn new(root: TreeLink) -> Self {
        let mut stack: Vec<TreeLink> = vec![root];
        let mut i = 0;
        while i < stack.len() {
            let left = stack[i].as_mut().unwrap().borrow_mut().left.clone();
            let right = stack[i].as_mut().unwrap().borrow_mut().right.clone();
            if left.is_some() {
                stack.push(left);
            }
            if right.is_some() {
                stack.push(right);
            }
            i += 1;
        }
        CBTInserter { stack }
    }

    fn insert(&mut self, v: i32) -> i32 {
        let link = tree!(v);
        let n = self.stack.len();
        self.stack.push(link.clone());
        let mut parent = self.stack[(n - 1) / 2].as_mut().unwrap().borrow_mut();
        if n % 2 == 1 {
            parent.left = link;
        } else {
            parent.right = link;
        }
        parent.val
    }

    fn get_root(&self) -> TreeLink {
        self.stack[0].clone()
    }
}

#[test]
fn test() {
    let mut obj = CBTInserter::new(tree!(1));
    assert_eq!(obj.insert(2), 1);
    assert_eq!(obj.get_root(), tree!(1, tree!(2), None));
}

Having problems with this solution? Click here to submit an issue on github.