Leetcode 99 Recover Binary Search Tree

Leetcode 99 Recover Binary Search Tree

The problem description is as follow:

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

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */

In case you are not familiar with Binary Search Tree, go through wiki from the link. A intuitive way is to get the in-order traversal. The correct order of in-order traversal of BST should be ascending. If two elements are swapped in BST, it will be reflected through in-order traversal’s order.

Naive Approach

Edit Recover Binary Search Tree on GitHub

1. Get the in-order traversal list of BST
2. Traverse the list, find the two elements in wrong order
3. Swap the two elements in wrong order

public class Solution {
    public void recoverTree(TreeNode root) {
        List<TreeNode> list = new LinkedList<TreeNode>();
        //get the inorder traveral list of tree
        builder(root, list);
        
        if(list.size()<2) {
            return;
        }
        
        TreeNode first = null;
        TreeNode second = null;
        
        //find two nodes with wrong order
        TreeNode pre = list.get(0);
        for(int i=1; i<list.size(); i++) {
            TreeNode cur = list.get(i);
            if(cur.val<pre.val) {
                if (first == null) {
                    first = pre;
                }
                second = cur;
            }
            pre = cur;
        }
        
        //swap two nodes back
        swap(first,second);
        return;
    }
    
    private void swap(TreeNode first, TreeNode second) {
        if(first!=null && second!=null){
            int temp = first.val;
            first.val = second.val;
            second.val = temp; 
        }
        return;
    }
    
    private void builder(TreeNode root, List<TreeNode> result) {
        if (root==null){
            return;
        }
        
        builder(root.left, result);
        result.add(root);
        builder(root.right, result);
        
    }
}

Building the in-order traversal list take O(n) time and O(log n) space. Finding the two misplaced elements take O(n) time and O(n) space since we are saving the whole list. Thus this algorithm takes O(n) time and space.

Better Approach

Edit Recover Binary Search Tree on GitHub

We can slightly improve the previous algorithm. Instead of saving the whole in-order traversal list, we could do the comparison when building the list.

public class Solution {
    TreeNode pre;
    TreeNode first;
    TreeNode second;
      
    private void inorder(TreeNode root){  
        if(root == null)  
            return;  

        inorder(root.left);  
        if(pre == null){  
            pre = root;  
        }else{  
            if(pre.val > root.val){  
                if(first == null){  
                    first = pre;
                }  
                second = root;  
            }  
            pre = root;  
        }  
        inorder(root.right);  
    }  
    
    private void swap(TreeNode first, TreeNode second) {
        if(first!=null && second!=null){
            int temp = first.val;
            first.val = second.val;
            second.val = temp; 
        }
        return;
    }
      
    public void recoverTree(TreeNode root) {  
        pre = null;  
        first = null;  
        second = null;  
        inorder(root);  
        swap(first,second);
    }
}

Building the in-order traversal list take O(n) time and O(log n) space. Two global variables saving the two misplaced elements take O(1) space. Thus this algorithm takes O(n) time and O(log n) space.

Morris Traversal Approach

Edit Recover Binary Search Tree on GitHub

After previous effort, we know that solving the problem heavily rely on in-order tree traversal. Since we can solve in-order binary tree traversal with O(n) time and O(1) space, we could further improve the algorithm by doing the comparison through morris traversal. This algorithm takes O(n) time and O(1) space.

public class Solution {
     private void swap(TreeNode first, TreeNode second) {
        if(first!=null && second!=null){
            int temp = first.val;
            first.val = second.val;
            second.val = temp; 
        }
        return;
    }
    
    public void recoverTree(TreeNode root) {
        TreeNode cur = root;
        TreeNode pre = null;
        TreeNode parent = null;
        TreeNode first = null;
        TreeNode second = null;
        boolean found = false;
        
        while(cur != null) {
            if(cur.left == null) {
                if(parent!=null && parent.val > cur.val) {
                    if(!found) {
                        first = parent;
                        found = true;
                    }
                    second = cur;
                }
                parent = cur;
                cur = cur.right;
            }
            else {
                pre = cur.left;
                while(pre.right!=null && pre.right!=cur)
                    pre = pre.right;
                if(pre.right==null) {
                    pre.right = cur;
                    cur = cur.left;
                }
                else {
                    pre.right = null;
                    
                    if(parent.val > cur.val) {
                        if(!found) {
                            first = parent;       
                            found = true;
                        }
                        second = cur;
                    }
                    parent = cur;
                    cur = cur.right;
                }
            }
        }
        
        swap(first,second);
    }
}