Chương trình sửa cây BST gần đúng thành BST chính xác trong Python

Cho cây nhị phân gần giống BST (Almost BST), chỉ có 2 nút bị hoán đổi giá trị. Sửa lại thành BST đúng.

Ví dụ: cây gốc inorder [2, 3, 6, 8, 9] (đã sửa).

Thuật toán Inorder Traversal: BST inorder cho dãy tăng dần. Nếu 2 nút bị swap, dãy inorder sẽ có 1 hoặc 2 cặp nghịch thế (inversion).

• Duyệt inorder, theo dõi prev_node
• Nếu node.val < prev_node.val → tìm thấy nghịch thế:
  • Nút đầu tiên của cặp nghịch thế đầu → max_node (lớn hơn, bị swap xuống)
  • Nút thứ hai của cặp nghịch thế cuối → min_node (nhỏ hơn, bị swap lên)
• Hoán đổi giá trị min_node và max_node

Triển khai

class TreeNode:
    def __init__(self, data, left=None, right=None):
        self.val = data
        self.left = left
        self.right = 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

class Solution:
    def solve(self, root):
        prev_node = None
        min_node = None
        max_node = None
        found_one = False
        
        for node in root:  # Sử dụng __iter__ inorder
            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)

# Inorder traversal để kiểm tra
def print_tree(root):
    if root:
        print_tree(root.left)
        print(root.val, end=', ')
        print_tree(root.right)

print_tree(ob.solve(root))

Đầu vào

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

Đầu ra

2, 3, 6, 8, 9,

Giải thích: Cây ban đầu inorder: [6, 3, 2, 8, 9] (2 cặp nghịch thế: 6>3 và 3>2). max_node=6 (từ cặp 6>3), min_node=2 (từ cặp 3>2). Swap 6↔2 → inorder: [2,3,6,8,9] đúng BST. Phức tạp O(n) time, O(h) space (stack đệ quy ẩn trong generator).