[파이썬 알고리즘] 24. 이진 트리 합의 범위

1. 문제

Leetcode 938. 938. Range Sum of BST

이진 탐색 트리가 주어졌을 때 low 이상, high 이하의 값을 가진 노드의 합을 구하여라.

 

2. 재귀 구조 DFS를 이용한 풀이

이진 탐색 트리의 정의는 지난 글에서 정리했고, 간단히 말하면 정렬된 트리라고 할 수 있다.

https://baechu-story.tistory.com/41 

[파이썬 알고리즘] 23. 이진 탐색 트리(BST)를 더 큰 수 합계 트리로

이 트리에 특정 범위에 있는 숫자들의 합을 구하기 위해 큰 수부터 작은 수 순서대로 탐색을 진행하는 방식으로 구현했다. BST는 더 큰 숫자가 오른쪽 노드에 있기 때문에, 오른쪽 노드, 자기 자신, 왼쪽 노드 순서로 탐색을 진행한다. 자기 자신을 탐색할 땐 자신의 값이 high와 low 사이면 자신의 값을 sum에 더했다. 만약 자신의 값이 low보다 작으면 node.left를 탐색할 필요가 없으므로 함수를 종료했다. 탐색 시 활용하는 sum, low, high는 모두 클래스 변수로 두고 활용했다.

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

 

 

3. BFS를 이용해 조건에 맞는 노드를 탐색하는 풀이

 책에서 제시된 풀이 중 하나는 스택을 만들고, 탐색할 노드를 스택에 넣고 마지막에 넣은 노드부터 탐색하는 BFS 방식이다. 스택의 제일 위에서 노드를 하나 꺼내고, 그 노드의 값이 low와 high 중간이면 노드의 값을 sum에 더한다. 만약 값이 low보다 크다면 왼쪽의 노드를 탐색해야 하기 때문에 node.left를 스택에 넣는다. 마찬가지로 값이 high보다 작다면 오른쪽의 노드를 탐색해야 하기 때문에 node.right를 스택에 넣는다. 스택에 노드가 없을 때까지 이 과정을 반복한다.

import collections
class Solution:
    def rangeSumBST(self, root: TreeNode, low: int, right: int) -> int:
        stack = collections.deque([root])
        sum = 0
        # 큐 연산을 이용해 반복 구조 BFS로 필요한 노드 탐색
        while stack:
            node = stack.popleft()
            if node:
                if node.val > low:
                    stack.append(node.left)
                if node.val < right:
                    stack.append(node.right)
                if low <= node.val <= right:
                    sum += node.val
        return sum