diff --git a/0623-add-one-row-to-tree/0623-add-one-row-to-tree.py b/0623-add-one-row-to-tree/0623-add-one-row-to-tree.py index 18a2406..f62ee51 100644 --- a/0623-add-one-row-to-tree/0623-add-one-row-to-tree.py +++ b/0623-add-one-row-to-tree/0623-add-one-row-to-tree.py @@ -10,16 +10,30 @@ def addOneRow(self, root: Optional[TreeNode], val: int, depth: int) -> Optional[ if depth == 1: return TreeNode(val, root, None) - elif depth == 2: - root.left = TreeNode(val, root.left, None) - root.right = TreeNode(val, None, root.right) - - else: - if root.left: - self.addOneRow(root.left, val, depth - 1) - if root.right: - self.addOneRow(root.right, val, depth - 1) + queue = [root] + while 0 < len(queue): + + depth -= 1 + + count = len(queue) + + for _ in range(count): + node = queue.pop(0) + + if node.left is not None: + queue.append(node.left) + + if node.right is not None: + queue.append(node.right) + + if depth == 1: + node.left = TreeNode(val, node.left, None) + node.right = TreeNode(val, None, node.right) + + if depth == 1: + break + return root