[LeetCode] 814 | Binary Tree Pruning | Python

---

문제 링크: https://leetcode.com/problems/binary-tree-pruning/

---

문제

Given the root of a binary tree, return the same tree where every subtree (of the given tree) not containing a 1 has been removed.

A subtree of a node node is node plus every node that is a descendant of node.

---

조건

• The number of nodes in the tree is in the range [1, 200].
• Node.val is either 0 or 1.

---

내 풀이

Tree 구조 관련 문제는 recursive하게 짠다.

dfs로 탐색하며 트리의 끝(node)에 도달했을 때, node.left와 node.right을 봤을 때 그 값이 0인 child는 없애야 하기 때문에 None으로 바꾼다.

node.val, node.left.val, node.right.val 중 하나라도 1이 있다면 그 노드는 없앨 수 없기 때문에 1로 리턴함으로써 그 값(1)을 그 위의 노드한테 전달한다.

---

코드

# 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 pruneTree(self, root: Optional[TreeNode]) -> Optional[TreeNode]:
        def search(root):
            if not root:
                return 0
            
            left = search(root.left)
            right = search(root.right)
            
            if not left:
                root.left = None
            if not right:
                root.right = None
            
            return max(root.val, left, right)
            
        if search(root):
            return root