Recover Binary Search Tree

recover.cpp

/**
 * Definition of TreeNode:
 * class TreeNode {
 * public:
 *     int val;
 *     TreeNode *left, *right;
 *     TreeNode(int val) {
 *         this->val = val;
 *         this->left = this->right = NULL;
 *     }
 * }
 */

class Solution {
public:
    /**
     * @param root: the given tree
     * @return: the tree after swapping
     */
     
     
    TreeNode *findKthNode(TreeNode *root, int k) {
        if (root == NULL) {
            return NULL;
        }
        TreeNode *res = NULL;
        stack<TreeNode *> stack;
        TreeNode *ptr = root;
        int count = 0;
        while(ptr != NULL || ! stack.empty()) {
            if (ptr != NULL) {
                stack.push(ptr);
                ptr = ptr->left;
            } else {
                ptr = stack.top();
                stack.pop();
                count++;
                if (count == k) {
                    res = ptr;
                    break;
                }
                ptr = ptr->right;
            }
        }
        return res;
    } 
    
     
    TreeNode * bstSwappedNode(TreeNode * root) {
        // write your code here
        
        if (root == NULL || (root->left == NULL && root->right == NULL)) {
            return root;
        }
        
        vector<int> values;
        stack<TreeNode *> stack;
        TreeNode *ptr = root;
        while(ptr != NULL || ! stack.empty()) {
            if (ptr!=NULL) {
                stack.push(ptr);
                ptr = ptr->left;
            } else {
                ptr = stack.top();
                values.push_back(ptr->val);
                stack.pop();
                ptr = ptr->right;
            }
        }
        
        int size = values.size();
        int left = 0, right = size-1;
        while(left < size) {
            if (values[left+1] < values[left]) {
                break;
            }
            left++;
        }

        while(right >= 0) {
            if (values[right-1] > values[right]) {
                break;
            }
            right--;
        }
        
        TreeNode *x = findKthNode(root, left+1);
        TreeNode *y = findKthNode(root, right+1);
        int t = x->val;
        x->val = y->val;
        y->val = t;
        return root;
    }
};