Binary Tree Pruning

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

Binary Tree Pruning - LeetCode

문제 해설 : 이진 트리의 root가 주어집니다. 같은 tree에서 1을 포함하지 않는 subtree를 제거하여 반환하세요.

 

제한 사항

1) tree의 node의 개수는 [1, 200]을 만족합니다.

2) Node.val은 0 또는 1입니다.

 

두 가지 방법으로 문제를 해결했습니다.

 

첫 번째 방법은, 트리를 전체 탐색하여 제거할 노드를 고르고 다시 트리를 순회하며 노드를 제거해나가는 방식입니다.

아래는 솔루션입니다.

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

class Solution:

    def search(self, node: TreeNode, nodes_to_be_removed: set) -> int:
        val = node.val

        if node.left:
            val += self.search(node.left, nodes_to_be_removed)
        if node.right:
            val += self.search(node.right, nodes_to_be_removed)

        if val == 0:
            nodes_to_be_removed.add(node)

        return val

    def remove(self, node: TreeNode, nodes_to_be_removed: set) -> None:
        if node.left:
            if nodes_to_be_removed.__contains__(node.left):
                node.left = None
            else:
                self.remove(node.left, nodes_to_be_removed)
        
        if node.right:
            if nodes_to_be_removed.__contains__(node.right):
                node.right = None
            else:
                self.remove(node.right, nodes_to_be_removed)

    def pruneTree(self, root: TreeNode) -> TreeNode:
        nodes_to_be_removed = set()
        self.search(root, nodes_to_be_removed)

        if nodes_to_be_removed.__contains__(root):
            return None
        
        self.remove(root, nodes_to_be_removed)
        return root

nodes_to_be_removed는 set객체로 제거할 노드를 저장합니다.

 

search 메서드는 이진 트리의 left와 right를 재귀적으로 탐색하며 val에 자식 노드의 값을 더해갑니다.

val에 저장된 값이 0이라면 자신을 포함한 자식노드의 합이 0이므로 제거대상이 됩니다.

따라서 val가 0인 노드를 nodes_to_be_removed에 추가합니다.

 

remove 메서드는 root부터 다시 전체 트리를 탐색하면서 자식 노드가 nodes_to_be_removed에 포함되어 있는지 확인합니다. 만약 포함되어 있다면 자식 노드와 연결을 끊습니다.

 

만약 root 노드 또한 nodes_to_be_removed에 포함되어 있다면 root 노드도 제거대상이므로

None을 반환합니다.

 

 

전체 트리를 2번에 걸쳐 순회하고 제거 대상 노드의 정보를 저장할 set 객체를 사용했는데

최적화 과정을 거치면 노드 정보를 저장하지 않고 트리를 한 번만 순회하여 문제를 해결할 수 있습니다.

 

# 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 search(self, node: TreeNode) -> int:
        val = node.val

        if node.left:
            leftVal = self.search(node.left)
            if leftVal == 0:
                node.left = None
            val += leftVal

        if node.right:
            rightVal = self.search(node.right)
            if rightVal == 0:
                node.right = None
            val += rightVal

        return val

    def pruneTree(self, root: TreeNode) -> TreeNode:
        self.search(root)
        if root.val == 0 and not root.left and not root.right:
            return None
        return root

root부터 전체 노드를 재귀적으로 탐색한다는 접근 방식은 동일합니다.

다만 제거 대상이 될 노드의 정보를 따로 저장하지 않고

자식 노드의 val가 0임을 확인하면 그 순간 연결을 끊도록 코드를 수정했습니다.

 

마지막으로 root의 값이 0이고 root가 자식 노드를 갖지 않는다면

root 노드 또한 제거 대상이 되므로 None을 반환하는 코드를 넣었습니다.

 

 

 

개선점

짧은 시간 내에 문제를 해결해야 하는 코딩테스트나 대회의 경우, 문제를 처음 보고 떠오른 아이디어를 코드로 바로 옮기는 능력이 중요하다고 생각합니다. 다만 채점이 끝난 후에는 더 나은 방법을 고민해봐야 합니다.

코딩 테스트는 내 실력을 남에게 보여주기 위한 수단일 뿐, 내가 가진 가장 높은 퍼포먼스를 증명하는 방법은 아니기 때문입니다. 더 나은 방법이 있는지 고민하고 개선할 점을 찾고 다른 사람과 아이디어를 비교해나간다면 조금씩 더 나아가는 개발자가 될 수 있을 거라 생각합니다.