Recover Binary Search Tree

LeetCode #99

Description:

Two elements of a binary search tree (BST) are swapped by mistake.

Recover the tree without changing its structure.

Note: A solution using O(n) space is pretty straight forward. Could you devise a constant space solution?

My Note: I am using recursion, which internally uses a stack. This method is clear and straight forward. I give up the O(1) space solution.

Example:

     6
   /  \
  4    9
 /\    /\
1  5  7  11

Inorder traversal: 1 4 5 6 7 9 11
Switch 4, 5: 1 5 4 6 7 9 11
Switch 4, 6: 1 6 5 4 7 9 11
Switch 5, 9: 1 4 9 6 7 5 11

Idea:

Traverse the tree Inorder, use three pointers:

1. pre, point to previously visited node
2. first, point to the first incorrect location
3. second, point to the second incorrect location

From the example above, it is easy to see, first should be the first of (a, b) pair which violate a<b property; second should be the second of (a, b) which violate the property

Code:

class Solution {
public:
    void recoverTree(TreeNode* root) {
        TreeNode *pre=NULL;
        TreeNode *first=NULL;
        TreeNode *second=NULL;

        locateSwapNodes(root, &pre, &first, &second);
        swap(first->val, second->val);
    }

    void locateSwapNodes(TreeNode *root, TreeNode **pre, TreeNode **first, TreeNode **second){
        if(root==NULL) return;

        locateSwapNodes(root->left, pre, first, second);

        if(*pre==NULL) 
            *pre=root;
        else{
            if((*pre)->val > root->val){
                if(*first==NULL){ 
                    *first=*pre;
                }
                *second=root;
            }
            *pre=root;
        }

        locateSwapNodes(root->right, pre, first, second);
    }
};