99. 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?

// Recursion with O(n) space
void helper(TreeNode* node, TreeNode*& pre, TreeNode*& first, TreeNode*& second) {
    if (!node) return;
    helper(node->left, pre, first, second);
    if (pre && pre->val >= node->val) {
        if (!first) first = pre;
        if (first) second = node;
    }
    pre = node;
    helper(node->right, pre, first, second);
}
void recoverTree(TreeNode* root) { // time: O(n); space: O(n)
    TreeNode *pre = nullptr, *first = nullptr, *second = nullptr;
    helper(root, pre, first, second);
    swap(first->val, second->val);
}

// Iteration with Stack
void recoverTree(TreeNode* root) { // time: O(n); space: O(n)
    TreeNode *pre = nullptr, *first = nullptr, *second = nullptr;
    stack<TreeNode*> st;
    TreeNode* cur = root;
    while (cur || !st.empty()) {
        while (cur) {
            st.push(cur);
            cur = cur->left;
        }
        cur = st.top(); st.pop();
        if (pre && pre->val > cur->val) {
            if (!first) first = pre;
            if (first) second = cur;
        }
        pre = cur;
        cur = cur->right;
    }
    swap(first->val, second->val);
}

// Morris Traversal
void recoverTree(TreeNode* root) { // time: O(n); space: O(1)
    TreeNode *parent = nullptr, *first = nullptr, *second = nullptr;
    TreeNode *pre = nullptr, *cur = root;
    while (cur) {
        if (!cur->left) {
            if (parent && parent->val >= cur->val) {
                if (!first) first = parent;
                if (first) 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;
                    if (first) second = cur;
                }
                parent = cur;
                cur = cur->right;
            }
        }
    }
    swap(first->val, second->val);
}