1382. Balance a Binary Search Tree

Balance a Binary Search Tree - LeetCode

이진 탐색 트리가 주어진다.

주어진 이진 탐색 트리를 균형 이진 탐색 트리로 만들어 반환하는 문제다.

균형 이진 탐색 트리는 모든 노드의 left, right subtree의 depth 차이가 1 이하인 이진 탐색 트리다.

 

1. 중위 순회를 이용한 트리 생성

class Solution:
    def balanceBST(self, root: TreeNode) -> TreeNode:

        seq = []

        def in_order(cur):
            if cur.left:
                in_order(cur.left)
            seq.append(cur.val)
            if cur.right:
                in_order(cur.right)

        in_order(root)

        def make_bst(candidate):
            mid = len(candidate) // 2
            cur_root = TreeNode(candidate[mid])
            
            if candidate[:mid]:
                cur_root.left = make_bst(candidate[:mid])
            if candidate[mid + 1:]:
                cur_root.right = make_bst(candidate[mid + 1:])

            return cur_root

        return make_bst(seq)

이진 탐색 트리가 주어지기 때문에 중위 순위를 하면 오름차순으로 노드를 조회할 수 있다.

 

균형 이진 탐색 트리이기 때문에 왼쪽 서브 트리, 오른쪽 서브 트리의 노드 개수를 최대 1개 차이만 나게 만들어주면 원하는 결과를 얻을 수 있다. 노드 후보 list를 들고 재귀에 진입해서 list의 중간 index에 있는 값을 root로 하고 index 왼쪽 노드로는 왼쪽 서브 트리를, 오른쪽 노드로는 오른쪽 서브 트리를 만드는 것을 반복하면 된다.

 

엄밀히 따지면 이건 기존 트리를 균형을 맞춰서 반환한 것이 아니라 균형을 맞는 새 트리를 만든 셈이다.

만일 공간 복잡도 제한이 걸려서 기존 트리를 가지고 균형을 맞춰야 한다면 자료구조 시간에 배웠던 AVL 트리 만들기 공식 같은 것을 사용해야 한다.

 

균형 잡힌 이진 탐색 트리: AVL 트리