Recover Binary Search Tree

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

Recover the tree without changing its structure.

Example 1:

Input: [1,3,null,null,2]

   1
  /
 3
  \
   2

Output: [3,1,null,null,2]

   3
  /
 1
  \
   2

Example 2:

Input: [3,1,4,null,null,2]

  3
 / \
1   4
   /
  2

Output: [2,1,4,null,null,3]

  2
 / \
1   4
   /
  3

Follow up:

• A solution using O(n) space is pretty straight forward.

• Could you devise a constant space solution?

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    void recoverTree(TreeNode* root) {
        TreeNode *first = NULL, *second = NULL, *parent = NULL;
        TreeNode *cur, *pre;
        cur = root;
        while (cur) {
            if (!cur->left) 
            {
                if (parent && parent->val > cur->val) 
                {
                    if (!first) 
                        first = parent;
                    second = cur;
                }
                parent = cur;
                cur = cur->right;
            } 
            else 
            {
                pre = cur->left;
                while (pre->right && pre->right != cur) 
                    pre = pre->right;
                if (!pre->right) 
                {
                    pre->right = cur;
                    cur = cur->left;
                } 
                else 
                {
                    pre->right = NULL;
                    if (parent->val > cur->val) 
                    {
                        if (!first) first = parent;
                        second = cur;
                    }
                    parent = cur;
                    cur = cur->right;
                }
            }
        }
        if (first && second) swap(first->val, second->val);
    }
};