We could keep all node whose left node or right node is empty in a queue and once inserting a new node, add it to the top of the queue's left or right.
# Definition for a binary tree node.
# class TreeNode:
# def __init__(self, x):
# self.val = x
# self.left = None
# self.right = None
from collections import deque
class CBTInserter:
def __init__(self, root: TreeNode):
self.root = root
self.queue = deque([root])
while self.queue[0].left and self.queue[0].right:
node = self.queue.popleft()
self.queue.append(node.left)
self.queue.append(node.right)
def insert(self, v: int) -> int:
node = TreeNode(v)
self.queue.append(node)
if self.queue[0].left is None:
self.queue[0].left = node
return self.queue[0].val
elif self.queue[0].right is None:
self.queue[0].right = node
return self.queue.popleft().val
def get_root(self) -> TreeNode:
return self.root
# Your CBTInserter object will be instantiated and called as such:
# obj = CBTInserter(root)
# param_1 = obj.insert(v)
# param_2 = obj.get_root()