Leetcode 955 Solution

This article provides solution to leetcode question 955 (complete-binary-tree-inserter)

https://leetcode.com/problems/complete-binary-tree-inserter

Solution

# 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):
        self.nodes = [[root]]

        i = 0
        while i < len(self.nodes):
            for node in self.nodes[i]:
                if (node.left or node.right) and i == len(self.nodes) - 1:
                    self.nodes.append([])
                if node.left:
                    self.nodes[i + 1].append(node.left)
                if node.right:
                    self.nodes[i + 1].append(node.right)
            i += 1

    def insert(self, v: int) -> int:
        if len(self.nodes) == 1 or len(self.nodes[-1]) == len(self.nodes[-2]) * 2:
            self.nodes.append([])

        parent = self.nodes[-2][len(self.nodes[-1]) // 2]

        if len(self.nodes[-1]) % 2 == 0:
            parent.left = TreeNode(v)
            self.nodes[-1].append(parent.left)
        else:
            parent.right = TreeNode(v)
            self.nodes[-1].append(parent.right)

        return parent.val

    def get_root(self) -> TreeNode:
        return self.nodes[0][0]


# Your CBTInserter object will be instantiated and called as such:
# obj = CBTInserter(root)
# param_1 = obj.insert(v)
# param_2 = obj.get_root()