Problem description
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
- Since the tree is complete binary tree, we can use serialize it to an array. Everytime, we insert we append to the end of the array and find the parent of the inserted node and modify the left and right attribute of the parent node.
Below is python implementation
# Definition for a binary tree node.
# class TreeNode:
# def __init__(self, x):
# self.val = x
# self.left = None
# self.right = None
class CBTInserter:
def __init__(self, root: TreeNode):
def count(root)->int:
if not root:
return 0
return 1 + count(root.left) + count(root.right)
n = count(root)
self.arr = [0 for i in range(n)]
def dfs(root, i)->None:
if not root:
return
self.arr[i] = root
dfs(root.left, 2*i + 1)
dfs(root.right, 2*i + 2)
return
dfs(root, 0)
def insert(self, v: int) -> int:
node = TreeNode(v)
n = len(self.arr)
parent = None
if n % 2 == 1:
parent = self.arr[n // 2]
parent.left = node
else:
parent = self.arr[(n - 1) // 2]
parent.right = node
self.arr.append(node)
return parent.val
def get_root(self) -> TreeNode:
return self.arr[0]
# Your CBTInserter object will be instantiated and called as such:
# obj = CBTInserter(root)
# param_1 = obj.insert(v)
# param_2 = obj.get_root()