Recover Binary Search Tree

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?
-------------------- codes ---------------------------------------
/**
 * Definition of TreeNode:
 * class TreeNode {
 * public:
 *     int val;
 *     TreeNode *left, *right;
 *     TreeNode(int val) {
 *         this->val = val;
 *         this->left = this->right = NULL;
 *     }
 * }
 * Example of iterate a tree:
 * Solution iterator = Solution(root);
 * while (iterator.hasNext()) {
 *    TreeNode * node = iterator.next();
 *    do something for node
 */
class Solution {
private:
    TreeNode *cur;
    bool hasItor = false;
public:
    //@param root: The root of binary tree.
    Solution(TreeNode *root) {
        // write your code here
        // update the right links
        cur = root;
    }

    //@return: True if there has next node, or false
    bool hasNext() {
        // write your code here
        TreeNode *pre = NULL;
        while (cur != NULL) {
            if (cur->left == NULL) {
                return true;
            } else {
                pre = cur->left;
                while (pre->right && pre->right != cur) {
                    pre = pre->right;
                }
                if (pre->right == NULL) {
                    pre->right = cur;
                    cur = cur->left;
                } else {
                    pre->right = NULL;
                    hasItor = true;
                    return true;
                }
            }
        }
        return false;
    }
    
    //@return: return next node
    TreeNode* next() {
        // write your code here
        TreeNode *result = cur;
        if (hasItor) {
            cur = cur->right;
            return result;
        } else {
            TreeNode *pre = NULL;
            while (cur != NULL) {
                if (cur->left == NULL) {
                    cur = cur->right;
                    return result;
                } else {
                    pre = cur->left;
                    while (pre->right && pre->right != cur) {
                        pre = pre->right;
                    }
                    if (pre->right == NULL) {
                        pre->right = cur;
                        cur = cur->right;
                    } else {
                        pre->right = NULL;
                        cur = cur->right;
                        hasItor = false;
                        return result;
                    }
                }
            }
            
        }
    }
};

---------------- thinking --------------------
1. learn Morris Traversal Algorithm 
2. how to find the two switched nodes: two situations : 
    2,1, 3,4,5  -> nodes contain 2 items
    4, 2, 3, 1, 5 > nodes contain 4 items
https://oj.leetcode.com/discuss/13034/no-fancy-algorithm-just-simple-and-powerful-order-traversal
http://www.cnblogs.com/AnnieKim/archive/2013/06/15/morristraversal.html ------------------ codes -----------------------
class Solution {
public:
    void recoverTree(TreeNode *root) {
        //morris traversal
        TreeNode *cur = root;
        TreeNode *pre_node = NULL;
        vector<TreeNode*> nodes;
        while (cur != NULL) {
            if (cur->left == NULL) {
                if (pre_node && cur->val < pre_node->val) {
                    nodes.push_back(pre_node);
                    nodes.push_back(cur);
                } 
                pre_node = cur;
                cur = cur->right;
            } else {
                TreeNode *pre = cur->left;
                while (pre->right != NULL && pre->right != cur) {
                    pre = pre->right;
                }
                if (pre->right == cur) {
                    pre->right = NULL;
                    if (pre_node && cur->val < pre_node->val) {
                        nodes.push_back(pre_node);
                        nodes.push_back(cur);
                    } 
                    pre_node = cur;
                    cur = cur->right;
                } else {
                    pre->right = cur;
                    cur = cur->left;
                }
            }
        }
        if (nodes.size() == 2) {
            int tmp = nodes[0]->val;
            nodes[0]->val = nodes[1]->val;
            nodes[1]->val = tmp;
        } else {
            int tmp = nodes[0]->val;
            nodes[0]->val = nodes[3]->val;
            nodes[3]->val = tmp;            
        }
    }

};