L52. Recover BST | Correct BST with two nodes swapped

Recovering a Binary Search Tree (BST) where exactly two nodes have been swapped requires restoring the BST property without violating optimal time and space boundaries. By leveraging the inherent property of in-order traversals on Binary Search Trees, we can identify and correct the misplaced nodes efficiently.

1. Why Do We Need This? (Problem Statement & Status Quo)

A Binary Search Tree (BST) is structured such that for any given node, all values in its left subtree are strictly smaller than the node, and all values in its right subtree are strictly greater than the node.

In this problem, you are given the root of a complete BST where exactly two nodes have been swapped by mistake. The structural topology of the tree remains a valid binary tree, but the BST ordering invariant is broken. Your goal is to identify those two swapped nodes and correct their values in-place so that the tree becomes a valid BST once again.

Incorrect BST (3 and 2 are swapped):
        3
       / \
      1   4
         /
        2

Corrected BST (2 and 3 swapped back):
        2
       / \
      1   4
         /
        3

2. Core Mental Model & Analogy

Real-Life Analogy: The Alphabetized File Cabinet

Imagine a massive file cabinet where folders are supposed to be arranged in strictly alphabetical order from left to right. Someone walks up, pulls out two folders, and accidentally places them in each other's slots.

If you walk down the row checking each folder against the one immediately preceding it, you will notice unexpected "dips" where a folder's name is alphabetically smaller than the folder that came directly before it. These dips are the violations of our sorted sequence. By tracking where these dips occur, you can pinpoint the exact two folders that were misplaced and swap them back.

3. The Naive Brute-Force Approach & Its Limitations

The most straightforward way to solve this problem is to flatten or capture the traversal, sort it, and overwrite the values.

Brute-Force Workflow

1. Perform any tree traversal (e.g., In-order, Pre-order, Post-order) to collect all node values into an array.
2. Sort the array using a standard sorting algorithm (which yields a correctly ordered sequence).
3. Perform a second traversal over the tree, assigning the sorted values back to the nodes sequentially.

Complexity Breakdown

• Time Complexity: $\mathcal{O}(N) + \mathcal{O}(N \log N) + \mathcal{O}(N) \approx \mathcal{O}(N \log N)$, dominated by the sorting step.
• Space Complexity: $\mathcal{O}(N)$ auxiliary space to store all node values in a list or array, plus recursive stack space.

Note: Interviewers will reject this approach because it incurs unnecessary $\mathcal{O}(N \log N)$ time and $\mathcal{O}(N)$ extra space, ignoring the structural guarantees of a BST.

4. Step-by-Step Evolution: The In-Order Traversal Insight

The breakthrough comes from remembering a fundamental property of Binary Search Trees: The in-order traversal of a valid BST always yields a sorted sequence in ascending order.

If two nodes in a BST are swapped, their placement corrupts this sorted in-order sequence. Therefore, if we perform an in-order traversal (Left $\rightarrow$ Node $\rightarrow$ Right), our sequence will exhibit violations where a current node's value is less than or equal to the previous node's value.

Depending on whether the two swapped nodes were adjacent or non-adjacent in the in-order sequence, the traversal will encounter either one or two violation points.

5. Analyzing the Two Possible Cases of Swapped Nodes

When two elements are swapped in a sorted sequence, they create disruptions in the natural ascending order.

6. Comprehensive Flow Walkthrough: Detecting Violations

We maintain pointer references during our in-order traversal:

• previous: Tracks the previously visited node in the in-order traversal.
• first: Stores the first node where a violation occurs.
• middle: Stores the second node of the first violation pair (crucial for the adjacent case).
• last: Stores the second node where a second violation occurs.

Violation Tracking Logic

if (previous != null && current.val < previous.val) {
    if (first == null) {
        // This is the first violation
        first = previous;
        middle = current;
    } else {
        // This is the second violation
        last = current;
    }
}
previous = current;

Flow Diagram for Non-Adjacent Swaps

Flow Diagram for Adjacent Swaps

7. Production-Ready Code Implementation

Below is the optimal implementation in C++ and Java. Both implementations achieve $\mathcal{O}(N)$ time complexity and $\mathcal{O}(H)$ auxiliary stack space (where $H$ is the height of the tree).

C++ Implementation

class Solution {
private:
    TreeNode* first;
    TreeNode* prev;
    TreeNode* middle;
    TreeNode* last;

    void inorder(TreeNode* root) {
        if (root == nullptr) return;

        // Traverse left subtree
        inorder(root->left);

        // Process current node (Business logic)
        if (prev != nullptr && root->val < prev->val) {
            // If this is the first violation, mark first and middle
            if (first == nullptr) {
                first = prev;
                middle = root;
            } 
            // If this is the second violation, mark last
            else {
                last = root;
            }
        }
        prev = root;

        // Traverse right subtree
        inorder(root->right);
    }

public:
    void recoverTree(TreeNode* root) {
        first = middle = last = nullptr;
        prev = new TreeNode(INT_MIN); // Initialize with a very small value
        
        inorder(root);

        // If both violations were found, swap non-adjacent nodes
        if (first && last) {
            swap(first->val, last->val);
        } 
        // If only one violation was found, swap adjacent nodes
        else if (first && middle) {
            swap(first->val, middle->val);
        }
    }
};

Java Implementation

class Solution {
    private TreeNode first;
    private TreeNode prev;
    private TreeNode middle;
    private TreeNode last;

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

        // Traverse left subtree
        inorder(root.left);

        // Process current node (Business logic)
        if (prev != null && root.val < prev.val) {
            // If this is the first violation, mark first and middle
            if (first == null) {
                first = prev;
                middle = root;
            } 
            // If this is the second violation, mark last
            else {
                last = root;
            }
        }
        prev = root;

        // Traverse right subtree
        inorder(root.right);
    }

    public void recoverTree(TreeNode root) {
        first = middle = last = null;
        prev = new TreeNode(Integer.MIN_VALUE); // Initialize with minimum value
        
        inorder(root);

        // If both violations were found, swap non-adjacent nodes
        if (first != null && last != null) {
            int temp = first.val;
            first.val = last.val;
            last.val = temp;
        } 
        // If only one violation was found, swap adjacent nodes
        else if (first != null && middle != null) {
            int temp = first.val;
            first.val = middle.val;
            middle.val = temp;
        }
    }
}

8. Summary and Key Takeaways

• BST In-Order Invariant: The in-order traversal of a valid Binary Search Tree always produces a strictly sorted ascending sequence.
• Violation Detection: Swapping two nodes in a BST creates either one violation (adjacent swap) or two violations (non-adjacent swap) during an in-order traversal.
• Pointer Tracking: Maintaining prev, first, middle, and last pointers allows us to capture the misplaced nodes in a single traversal pass.
• Optimal Time Complexity: $\mathcal{O}(N)$ time is achieved by visiting each node exactly once during the recursive in-order traversal.
• Optimal Space Complexity: $\mathcal{O}(H)$ stack space (where $H$ is tree height), which can be further optimized to $\mathcal{O}(1)$ auxiliary space using Morris Traversal if required by strict space constraints.
• In-Place Modification: Values are swapped directly within the existing tree nodes without allocating new data structures.