Leetcode problem: Recover the binary tree that has had two nodes swapped/

recover_binary_tree.cpp

/* This is an easy problem, but made difficult because you can only use O(1) space
 * Traverse the tree use inorder. Its easy to spot nodes that are out of place: They will not be in increasing order
 * In order to accomplish this, we need to do an in order traverse and keep track of the previous node
 * The first time previous is larger than the current node (root), both previous and root must be added because both could
 * potentially be in violation. If however we find another violation (i.e. previous is greater than current), then we can replace
 * the last value of the node array with the current (root) node. Why? Because two nodes have been swapped. We find the larger of
 * the two first (because of the in order traversal), so the next node must be smaller of the two.
 * 
 * Space: O(1) - only an array of size 2
 * Time Complexity: O(n), where n is the number of nodes (i.e. complexity of in-order traversal)
 */

class Solution {
public:

    TreeNode* findNodes(TreeNode* root, TreeNode* prev, TreeNode* nodes[2]) {
        if (!root) {
            return prev;
        }
        
        TreeNode* previous = findNodes(root->left, prev, nodes);

        if (previous && previous->val > root->val) {
            if (!nodes[0] && !nodes[1]) {
                nodes[0] = previous;
                nodes[1] = root;
            } else {
                nodes[1] = root;
            }
        }
        
        return findNodes(root->right, root, nodes); 
    }
    
    void recoverTree(TreeNode *root) {
        TreeNode* nodes[2] = {NULL, NULL};
        findNodes(root, NULL, nodes);
        
        int temp = nodes[0]->val;
        nodes[0]->val = nodes[1]->val;
        nodes[1]->val = temp;
    }
};