[LeetCode] 623. Add One Row to Tree

623. Add One Row to Tree

Medium

Given the root of a binary tree and two integers val and depth, add a row of nodes with value val at the given depth depth.

Note that the root node is at depth 1.

The adding rule is:

• Given the integer depth, for each not null tree node cur at the depth depth - 1, create two tree nodes with value val as cur's left subtree root and right subtree root.
• cur's original left subtree should be the left subtree of the new left subtree root.
• cur's original right subtree should be the right subtree of the new right subtree root.
• If depth == 1 that means there is no depth depth - 1 at all, then create a tree node with value val as the new root of the whole original tree, and the original tree is the new root's left subtree.

 

Example 1:

Input: root = [4,2,6,3,1,5], val = 1, depth = 2
Output: [4,1,1,2,null,null,6,3,1,5]

Example 2:

Input: root = [4,2,null,3,1], val = 1, depth = 3
Output: [4,2,null,1,1,3,null,null,1]

 

Constraints:

• The number of nodes in the tree is in the range [1, 104].
• The depth of the tree is in the range [1, 104].
• -100 <= Node.val <= 100
• -105 <= val <= 105
• 1 <= depth <= the depth of tree + 1

---

문제 풀이

• 원하는 깊이에 있는 노드에 새로운 자식들을 추가하는 문제다.
• 몇가지 조건이 있는데, 
• 주어진 깊이에 노드가 존재하지 않으면 새로운 2 개의 노드를 생성해서 연결해야한다. 
• 그리고 깊이가 1 이라면 현재 존재하는 트리는, 새로 추가될 root 노드의 왼쪽 자식이 되야한다.
• 나머지 규칙은 추가될 깊이의 노드에서 왼쪽, 오른쪽 자식을 그에 맞게 넣는다는 조건이다.
• 노드의 개수가 최악의 경우 10^4, 깊이의 최악의 경우가 10^4 이기 때문에 이진탐색으로 풀어도 문제가 없다.

소스 코드

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def addOneRow(self, root: Optional[TreeNode], val: int, depth: int) -> Optional[TreeNode]:
        def rec(node, d):
            if not node:
                return
                
            if depth - 1 == d:
                l_temp = node.left
                l_new = TreeNode(val)
                node.left = l_new
                l_new.left = l_temp
                
                r_temp = node.right
                r_new = TreeNode(val)
                node.right = r_new
                r_new.right = r_temp
                return 
            
            rec(node.left, d + 1)
            rec(node.right, d + 1)
        if depth == 1:
            new = TreeNode(val)
            new.left = root
            root = new
        else:
            rec(root, 1)
        return root