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 noderoot
;CBTInserter.insert(int v)
will insert aTreeNode
into the tree with valuenode.val = v
so that the tree remains complete, and returns the value of the parent of the insertedTreeNode
;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:
- The initial given tree is complete and contains between
1
and1000
nodes. CBTInserter.insert
is called at most10000
times per test case.- Every value of a given or inserted node is between
0
and5000
.
Rust Solution
use rustgym_util::*;
struct CBTInserter {
stack: Vec<TreeLink>,
}
impl CBTInserter {
fn new(root: TreeLink) -> Self {
let mut stack: Vec<TreeLink> = vec![];
stack.push(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.