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

1. 문제

leetcode 1038. Binary Search Tree to Greater Sum Tree

이진 탐색 트리(BST)의 루트가 주어지면, 해당 트리의 값을 트리에서 자기 자신보다 큰 수의 합 + 자기 자신의 값으로 바꾸어라.

 

2. 풀이 과정

먼저 이진 탐색 트리는 다음과 같은 특성을 가지고 있다.

 

1. 왼쪽의 서브트리에 노드 값이 자신 이하이다.

2. 오른쪽의 서브트리에 노드 값이 자신 이상이다.

3. 왼쪽과 오른쪽의 서브트리가 모두 이진 탐색 트리이다.

이처럼 이진 탐색 트리는 정렬된 트리라고 할 수 있고, 탐색 시 시간 복잡도가 O(logn)에 불과하다.

 

주어진 문제에서는 트리를 큰 값부터 차례대로 탐색해야 한다. 따라서 재귀함수를 이용해 오른쪽 노드를 먼저 탐색하고, 자기 자신, 왼쪽 노드 순서로 탐색을 진행한다. 탐색 과정에서 sum이라는 변수를 dfs 함수의 인자로 전달하고, 자신의 노드 값을 더해서 리턴해 지금까지 탐색한 노드의 값을 전달한다. 그리고 만약 현재 노드가 존재하지 않는다면 입력받은 sum을 그대로 리턴하고 함수를 종료한다.

 

3. 전체 코드

 

def bstToGst(self, root: TreeNode) -> TreeNode:
    def dfs(node, sum):
        if node:
            sum = dfs(node.right, sum)
            sum += node.val
            node.val = sum
            sum = dfs(node.left, sum)
        return sum

    dfs(root, 0)

    return root