99. Recover Binary Search Tree

Description

The recoverTree function corrects a binary search tree (BST) where exactly two nodes have been swapped by mistake. The goal is to identify these two nodes and swap them back to restore the BST property, where the left subtree of any node contains only nodes with values less than the node’s value, and the right subtree only nodes with values greater.

Intuition

The problem can be approached by performing an in-order traversal of the BST. In a correctly ordered BST, the values retrieved during an in-order traversal should appear in ascending order. When two nodes are swapped, this order will be violated in exactly two places. By identifying these two violations, we can determine which nodes need to be swapped back to correct the tree.

Approach

1. In-order Traversal: - Perform an in-order traversal of the tree. During this traversal, keep track of the previous node (prev) visited.
2. Identify Swapped Nodes: - As you traverse, check if the value of the current node is less than the value of the previous node (prev). If so, a violation of the BST property has been detected. - The first time a violation is found, mark the previous node as big and the current node as small. - The second time a violation is found, update the small node to the current node.
3. Swap Values: - Once the traversal is complete, the two identified nodes (big and small) are swapped to restore the BST property.
4. Edge Cases: - If the tree is empty or contains only one node, no action is needed.

Complexity

Time Complexity:

The time complexity is \(O(n)\), where n is the number of nodes in the tree. This is because each node is visited exactly once during the in-order traversal.

Space Complexity:

The space complexity is \(O(h)\), where h is the height of the tree. This space is used by the recursion stack during the in-order traversal. In the worst case, for a completely unbalanced tree, this could be (O(n)), while for a balanced tree, it would be (O(\log n)).

Code

C++

class Solution {
public:
    void recoverTree(TreeNode* root) {
        TreeNode *small = nullptr, *big = nullptr, *prev = nullptr;

        std::function inorder = [&](TreeNode* r) {
            if (!r) return;
            inorder(r->left);

            if (prev && prev->val > r->val) {
                small = r;
                if (big) return;
                big = prev;
            }

            prev = r;
            inorder(r->right);
        };

        inorder(root);
        std::swap(small->val, big->val);
    }
};

Python

class Solution:
    def recoverTree(self, root: Optional[TreeNode]) -> None:
        small = big = prev = None

        def inorder(r):
            nonlocal small, big, prev
            if not r:
                return
            inorder(r.left)
            if prev and prev.val > r.val:
                small = r
                if not big:
                    big = prev
            prev = r
            inorder(r.right)

        inorder(root)
        small.val, big.val = big.val, small.val

Java

class Solution {
    private TreeNode small = null, big = null, prev = null;

    public void recoverTree(TreeNode root) {
        inorder(root);
        int temp = small.val;
        small.val = big.val;
        big.val = temp;
    }

    private void inorder(TreeNode r) {
        if (r == null) return;

        inorder(r.left);

        if (prev != null && prev.val > r.val) {
            small = r;
            if (big != null) return;
            big = prev;
        }

        prev = r;
        inorder(r.right);
    }
}

JavaScript

var recoverTree = function (root) {

    let small = null, big = null, prev = null;
    let inorder = function (r) {

        if (r == null) return;
        inorder(r.left);
        if (prev && prev.val > r.val) {

            small = r;
            if (big) return;
            big = prev;
        }
        prev = r;
        inorder(r.right);
        return;
    }
    
    inorder(root);
    [small.val, big.val] = [big.val, small.val];
};