BST iterator with early termination


  • 0
    I

    // We iterator through the BST using an interator.
    // We terminate as soon as both swapped nodes are identified.

    public void recoverTree(TreeNode root) {
    
        TreeIterator iter = new TreeIterator(root);
        
        TreeNode current = null;
        TreeNode previous = null;
        TreeNode firstWrong = null;
        
        while (true)
        {
            current = iter.getNext();
            
            if (firstWrong == null && previous != null && previous.val > current.val)
            {
                firstWrong = previous;    
            }
            else if (current == null || (firstWrong != null && current.val > firstWrong.val))
            {
                int temp = firstWrong.val;
                firstWrong.val = previous.val;
                previous.val = temp;
                break;
            }
            
            previous = current;
        }
    }
    
    public class TreeIterator 
    {
        Deque<TreeNode> stack;
                   
        public TreeNode getNext()
        {
            TreeNode popped = stack.pollLast();
            if (popped!=null)pushLeft(popped.right);
            return popped;
        }
        
        public TreeIterator (TreeNode root)
        {
            stack = new ArrayDeque<>();
            pushLeft(root);
        }
        
        private void pushLeft(TreeNode root)
        {
            TreeNode dummy = new TreeNode(0);
            dummy.left = root;
            while (dummy.left != null) stack.offerLast(dummy = dummy.left);
        }
    }

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.