LeetCode Solutions Blog

Recover Binary Search Tree

Number: 99

Difficulty: Medium

Paid? No

Problem Description

Given the root of a binary search tree (BST) where exactly two nodes have been swapped by mistake, restore the BST without changing its structure.

Key Insights

In a valid BST, an in-order traversal produces a sorted sequence.
A swapped node pair causes a deviation from the sorted order. There can be one or two such violations.
By tracking the previous node during an in-order traversal, we can identify nodes where the value order is incorrect.
After finding the two nodes, swap their values to fix the tree.
While the straightforward recursive in-order traversal uses O(n) space (due to the recursion stack), the Morris traversal technique can achieve O(1) space.

Space and Time Complexity

Time Complexity: O(n)
Space Complexity: O(n) in the recursive approach (or O(1) using Morris Traversal)

Solution

The solution involves performing an in-order traversal of the BST to detect the two nodes that have been swapped. During the traversal, we maintain a pointer to the previously visited node. When we encounter a node that is less than its previous node, we record the two nodes that violate the BST invariant. In cases where the swapped nodes are not adjacent, we will see two such violations; otherwise, we see only one. Finally, we swap the values of the two incorrect nodes, thereby restoring the BST.

Code Solutions

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right

def recoverTree(root):
    # Initialize variables to track the first and second swapped nodes and the previous node in in-order traversal
    first = second = prev = None

    # Helper function to perform in-order traversal
    def inorder(node):
        nonlocal first, second, prev
        if not node:
            return
        # Traverse the left subtree
        inorder(node.left)
        # Check for swapped nodes: if previous node's value is greater than the current node's value, violation occurs
        if prev and prev.val > node.val:
            if not first:
                # Record the first element when the first violation is encountered
                first = prev
            # Update the second element in every violation
            second = node
        # Update the previous node to the current node
        prev = node
        # Traverse the right subtree
        inorder(node.right)
    
    inorder(root)
    # Swap the values of the two nodes to recover the BST
    first.val, second.val = second.val, first.val