LeetCode: Recover Binary Search Tree

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?

---

We can perform a inorder traversal to find those two nodes with wrong values and then swapped them.

C++ Code:

/*
 * func: recover_BST
 * goal: recover binary search tree
 * @param root: root node of the tree
 * return:
 */
void recover_BST(TreeNode *root){
    TreeNode *iter = root;
    TreeNode *f1, *f2;
    bool found = false;
    TreeNode *previous = nullptr;
    
    while(iter){
        if(iter->left == nullptr){
            if(previous && previous->val > iter->val){
                if(!found){
                    f1 = previous;
                    found = true;
                }
                f2 = iter;
            }
            previous = iter;
            iter = iter->right;
        }else{
            TreeNode *pre = iter->left;
            while(pre->right && pre->right != iter)
                pre = pre->right;
            
            if(pre->right == nullptr){
                pre->right = iter;
                iter = iter->left;
            }
            else{
                pre->right = nullptr;
                if(pre->val > iter->val){
                    if(!found){
                        f1 = pre;
                        found = true;
                    }
                    f2 = iter;
                }
                previous = iter;
                iter = iter->right;
            }
        }
    }
    if(found){
        std::swap(f1->val, f2->val);
    }
}

Python Code:

# func: recover a BST
# @param root: root node of the linked list
# @return: root node
def recover_BST(root):
    iter = root
    found = False
    previous = None
    while iter:
        if not iter.left:
            if previous and previous.val > iter.val:
                if not found:
                    f1 = previous
                    found = True
                f2 = iter

            previous = iter
            iter = iter.right
        else:
            pre = iter.left
            while pre.right and pre.right != iter:
                pre = pre.right

            if not pre.right:
                pre.right = iter
                iter = iter.left
            else:
                pre.right = None
                if pre.val > iter.val:
                    if not found:
                        f1 = pre
                        found = True
                    f2 = iter
                previous = iter
                iter = iter.right
    if found:
        tmp = f1.val
        f1.val = f2.val
        f2.val = tmp
    return root