Complete Binary Tree Inserter

LeetCode Q 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:

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

Solution:

First, use O(n) time to build a nodesMap, where key is the index and value is the TreeNode.
When ever we want to insert a node, we can just use O(1) time to lookup that table find its parent and then insert it.

Code:

Map<Integer, TreeNode> nodesMap;
int index = 0;
public CBTInserter(TreeNode root) {
  nodesMap = new HashMap<>();
  buildsMap(root);
}

private void buildsMap(TreeNode root) {
  Queue<TreeNode> que = new LinkedList<>();
  que.offer(root);
  
  while (!que.isEmpty()) {
    TreeNode curr = que.poll();
    nodesMap.put(++index, curr);
    if (curr.left != null) que.offer(curr.left);
    if (curr.right != null) que.offer(curr.right);
  }
}

public int insert(int v) {
  TreeNode n = new TreeNode(v);
  TreeNode p = nodesMap.get((++index) / 2);
  
  if (p.left == null) p.left = n;
  else p.right = n;
  
  nodesMap.put(index, n);
  
  return p.val;
}

public TreeNode get_root() {
  return nodesMap.get(1);    
}