파이썬으로 거의 BST인 이진 트리를 정확한 BST로 복구하는 방법

어떤 이진 트리가 있고, 이 트리가 거의 이진 탐색 트리(BST)라고 가정해 봅시다. 즉, 단 두 개의 노드 값만 서로 바뀌어 있는 상태입니다. 우리는 이 트리를 올바르게 수정하여 정상적인 이진 탐색 트리를 반환해야 합니다.

예를 들어 입력이 다음과 같다면,

출력은 다음과 같아야 합니다.

문제 해결 접근 방식

이 문제를 해결하기 위해 다음 단계를 따릅니다.

• 초기화: prev_node := null, min_node := null, max_node := null
• 플래그 설정: found_one := False
• root의 중위 순회(inorder traversal) 결과에서 각 노드에 대해 다음을 수행합니다.
  • prev_node가 null이 아니라면:
    • 현재 노드의 값 < prev_node의 값이라면:
      • min_node가 null이거나 현재 노드의 값 < min_node의 값이면, min_node := node로 갱신합니다.
      • max_node가 null이거나 max_node의 값 < prev_node의 값이면, max_node := prev_node로 갱신합니다.
      • found_one이 이미 참이라면 루프를 종료(break)합니다.
      • 그렇지 않으면 found_one := True로 설정합니다.
    • prev_node := node로 갱신합니다.
• min_node와 max_node의 값을 서로 교환(swap)합니다.
• root를 반환합니다.

핵심 아이디어는 BST의 중위 순회 결과가 항상 오름차순으로 정렬되어 있다는 성질을 활용하는 것입니다. 순회 과정에서 값이 감소하는 지점을 발견하면 그 지점들이 바뀐 두 노드와 관련이 있으므로, 최솟값 후보(min_node)와 최댓값 후보(max_node)를 추적한 뒤 마지막에 두 값을 맞바꾸면 됩니다.

더 잘 이해하기 위해 다음 구현을 살펴보겠습니다.

예제 코드

class TreeNode:
    def __init__(self, data, left = None, right = None):
        self.val = data
        self.left = left
        self.right = right
    
def print_tree(root):
    if root is not None:
        print_tree(root.left)
        print(root.val, end = ', ')
        print_tree(root.right)
    
def __iter__(self):
    if self.left:
        for node in self.left:
            yield node
    yield self
    if self.right:
        for node in self.right:
            yield node

setattr(TreeNode, "__iter__", __iter__)
class Solution:
    def solve(self, root):
        prev_node = None
        min_node = None
        max_node = None
        found_one = False
        for node in root:
            if prev_node:
                if node.val < prev_node.val:
                    if min_node is None or node.val < min_node.val:
                        min_node = node
                    if max_node is None or max_node.val < prev_node.val:
                        max_node = prev_node
                    if found_one:
                        break
                    else:
                        found_one = True
            prev_node = node
        min_node.val, max_node.val = max_node.val, min_node.val
        return root
        
ob = Solution()
root = TreeNode(3)
root.left = TreeNode(6)
root.right = TreeNode(8)
root.right.left = TreeNode(2)
root.right.right = TreeNode(9)
print_tree(ob.solve(root))

입력

root = TreeNode(3)
root.left = TreeNode(6)
root.right = TreeNode(8)
root.right.left = TreeNode(2)
root.right.right = TreeNode(9)

출력

2, 3, 6, 8, 9,

출력 결과를 보면 중위 순회 값이 2, 3, 6, 8, 9로 오름차순 정렬된 것을 확인할 수 있습니다. 이는 트리가 올바른 이진 탐색 트리로 성공적으로 복구되었음을 의미합니다. 이 알고리즘은 중위 순회를 한 번만 수행하므로 시간 복잡도는 O(n)이며, 재귀 호출 스택 또는 반복자에 의한 공간 복잡도는 트리의 높이에 비례합니다.