Sunday, September 28, 2014

Recover Two Swapped Nodes in BST

[Question]
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?

[Analysis]
In-order traversal will reveal the swapped elements. If output the traversal result to an array, the space will be O(N). If using recursive calls, the space will be O(lg N). An in-place traversal will do it.

[Solution]
/////////////////
//Solution 1: Using in-place space
//
void inOrderTraversal(TreeNode* root) {
    TreeNode* pre = NULL;
    TreeNode* p = root;
    while(p! = NULL) {
            while (p->left!=NULL && p->left!=pre) {
                     TreeNode * q = p;
                     while (q->right!=NULL) q = q->right;
                     q->right = pre;
                     pre = p;
                     p = p->left;
            }
            cout<<p->val<<endl;
            pre = p;
            p = p->right;
            if (p->left == pre) {
                pre->right = NULL;
            }
    }

}


///////////////
// Solution 2: Recursive using a stack.
//
/**
 * Definition for binary tree
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */

    void recoverTree(TreeNode *root) {

        if (!root) return;
        stack<TreeNode*> st;
        TreeNode* cur=NULL;
        TreeNode* p1=NULL, *p2=NULL;
     
        st.push(root);
     
        while (!st.empty()) {
            while (st.top()->left) {
                st.push(st.top()->left);
            }

            if (!st.empty() && cur && cur->val > st.top()->val) {
                if (!p1) { p1 = cur; p2 = st.top();}
                else p2 = st.top();
            }

            while (!st.empty()&&st.top()->right==NULL) {
                cur = st.top();
                st.pop();
                if (!st.empty() && cur && cur->val > st.top()->val) {
                    if (!p1) { p1 = cur; p2 = st.top();}
                    else p2 = st.top();
                }
            }
         
            if (!st.empty()) {
                cur = st.top();
                st.pop();
                st.push(cur->right);
            }
        }
        if (p1 && p2 ) swap( p1->val, p2->val);
    }