[파이썬 알고리즘] 14장 - BST 합의 범위 (leetcode 938) - with pruning

leetcode.com/problems/range-sum-of-bst/

Range Sum of BST - LeetCode

 

<문제>

Given the root node of a binary search tree,

return the sum of values of all nodes with a value in the range [low, high].

 

<예시>

Input: root = [10,5,15,3,7,null,18], low = 7, high = 15

Output: 32

 

---

 

처음에는 성능 생각하지 않고 그냥 모든 노드를 방문하는 식의 비효율적인 코드를 작성했다.

(아래와 유사하게 코드를 작성하였다.)

class BookSolution1:
    def rangeSumBST(self, root: TreeNode, low: int, high: int) -> int:
        if not root:
            return 0

        return (root.val if low <= root.val <= high else 0) + \
                self.rangeSumBST(root.left, low, high) + \
                self.rangeSumBST(root.right, low, high)

 

하지만 BST의 정의를 이용하면 tree에 대해 pruning을 진행하면 속도를 매우 향상시킬 수 있다.

즉, 불필요한 node에 대해 탐색을 진행하지 않음으로써 효율적인 탐색을 진행하는 것이다.

class BookSolution2:
    def rangeSumBST(self, root: TreeNode, low: int, high: int) -> int:
        def dfs(node: TreeNode):
            if not node:
                return 0
            if node.val < low:
                # 현재 노드의 value가 low보다 작다면
                # left subtree에 존재하는 모든 node의 value는 low보다 작기 떄문에
                # right subtree만 탐색
                return dfs(node.right)
            elif node.val > high:
                # 현재 노드의 value가 high보다 크다면
                # right subtree에 존재하는 모든 node의 value는 high보다 크기 떄문에
                # left subtree만 탐색
                return dfs(node.left)
            else:
                return node.val+dfs(node.left)+dfs(node.right)
        return dfs(root)

pruning을 진행한 코드는 위와 같은데,

보면 현재 노드의 value가 low보다 작다면 왼쪽 subtree는 탐색하지 않고,

현재 노드의 value가 high보다 크다면 오른쪽 subtree를 탐색하지 않는다.

 

BST의 정의에 따르면 특정 노드의 왼쪽 subtree에는 자기자신보다 작은 값들만 위치하고, 오른쪽 subtree에는 자기자신보다 큰 값들만 위치하기 때문에 위와 같은 코드를 작성할 수 있는 것이다.

 

첫 번째 코드는 runtime: 280ms (faster than 17%) 의 성능이 나온 반면,

pruning을 진행한 두 번째 코드는 runtime: 192ms (faster than 97%)가 나왔다.

 

이 문제에서 얻어갈 수 있는 소중한 아이디어 같다 .. !