Two elements of a binary search tree (BST) are swapped by mistake. Recover the tree without changing its structure.

recover-bst-on.cc

/**
 * 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 {
    TreeNode* fst;
    TreeNode* sec;
    
public:
    void recoverTree(TreeNode* root) {
        dfs(root, NULL);
        
        int tmp = fst->val;
        fst->val = sec->val;
        sec->val = tmp;
    }
    
    TreeNode *dfs(TreeNode *root, TreeNode *last) {
        if (root->left != NULL) {
            last = dfs(root->left, last);
        }
        
        if (last != NULL && last->val > root->val) {
            if (!fst) {
                fst = last;
                sec = root;
            } else {
                sec = root;
                return NULL;
            }
        }
        
        if (root->right != NULL) {
            return dfs(root->right, root);
        }
        
        return root;
    }
};