[LeetCode/Python] Range Sum Of BST

Problem

Given the root node of a binary search tree and two integers low and high, return the sum of values of all nodes with a value in the inclusive range [low, high].

Example 1:

Input: root = [10,5,15,3,7,null,18], low = 7, high = 15
Output: 32
Explanation: Nodes 7, 10, and 15 are in the range [7, 15]. 7 + 10 + 15 = 32.

• 주어진 이진 탐색 트리의 값 중 범위 [low, high]에 속하는 값들의 합을 구하는 문제이다.

Solution

처음에는 단순하게 현재 노드의 값에 따라 조건문을 통해 범위에 맞는 노드만 탐색하는 코드를 작성했다.

class Solution:
    def rangeSumBST(self, root: Optional[TreeNode], low: int, high: int) -> int:
        answer = 0
        def visit(node: Optional[TreeNode]):
            if not node:
                return
            if low <= node.val <= high:
                nonlocal answer
                answer += node.val
            if low == node.val:
                visit(node.right)
            elif high == node.val:
                visit(node.left)
            else:
                visit(node.left)
                visit(node.right)
        visit(root)
        return answer

• 재귀함수로 각 노드를 탐색하는 방식이다.
• 현재 노드의 값이 범위에 포함된 값이라면 answer에 더해준다.
• 현재 노드의 값이 low와 같다면 오른쪽 노드만 탐색하고, high와 같다면 왼쪽 노드만 탐색한다.
• 둘 다 아니라면 왼쪽과 오른쪽 노드를 모두 탐색한다.

코드는 잘 동작하지만 조건문을 줄여서 최적화를 하는 코드도 작성해보았다.

class Solution:
    def rangeSumBST(self, root: Optional[TreeNode], low: int, high: int) -> int:
        answer = 0
        def dfs(node: Optional[TreeNode]) -> int:
            if not node:
                return 0
            current_val = 0
            if low <= node.val <= high:
                current_val = node.val
            left_sum = dfs(node.left)
            right_sum = dfs(node.right)
            return current_val + left_sum + right_sum
        return dfs(root)

• root 노드를 기준으로 dfs 알고리즘을 이용해 노드를 탐색한다.
• 탐색하려는 노드가 null이라면 탐색을 종료하고 0을 return 한다.
• 현재 노드의 값이 범위에 포함된 값이라면 current_val의 값을 노드의 val 값으로 업데이트한다.
• 현재 노드의 왼쪽 노드와 오른쪽 노드를 탐색해 구한 합을 current_val에 더하여 반환한다.