[Java] Add 3-line pruning to classic in-order traversal solution


  • -1
    S

    This solution is based on the most voted solution, and I add pruning to it so that we don't need traverse the whole tree for some cases.

            // pruning
            if (second != null && cur.val > first.val) {
                return;
            }
            visit(root.right);
    

    The idea is that when we visit cur (after we have set first and second) and find cur.val > first.val, we are sure second will not be updated anymore. So, we can return immediately without visiting root's right subtree.

    public class Solution {
    
        private TreeNode first = null;  // first wrong node
        private TreeNode second = null; // second wrong node
        private TreeNode prev = null;   // previous node before current node
    
        public void recoverTree(TreeNode root) {
            inorder(root);
            if (first != null && second != null) {
                int temp = first.val;
                first.val = second.val;
                second.val = temp;
            }
        }
    
        private void inorder(TreeNode root) {
            if (root == null) {
                return;
            }
            inorder(root.left);
            if (first == null && prev != null && prev.val > root.val) {
                first = prev;
            }
            if (first != null && prev.val > root.val) {
                second = root;
            }
            // pruning
            if (second != null && root.val > first.val) {
                return;
            }
            prev = root;
            inorder(root.right);
        }
    
    }
    

    Why?

    /**   
                                  swap first with s'
                              |                              |
            |                 |                              |       | 
            |  ...   |  ....  |   ..     ==>>      ..  | ..  | ..    |  
            |        |        |       |         |      |     |       |
         first       s       cur      s'        s'     s    cur      f
      */
    

    Suppose there is s' after cur which should be the real second, and then the in-order traversal array of tree should be sorted after we swap first and this s'.

    However, after swap, firstwill be after cur node, and then the interval [cur : f] is not sorted, therefore the assumption is wrong, so there is no possible second after cur.

    Example: in-order array of input tree: [1, 5, 3, 4, 2, 6, 7, 8, 9, 10] with root = 2, and correct recovery should be [1, 2, 3, 4, 5, 6, 7, 8, 9, 10].

    Before visit node(6), we will set first = node(5) and second = node (2), when visiting node(6), we know node(6) > first = 5, so we don't need check [7, 8, 9, 10] anymore. Because if we swap node(2) and node(5), job is done!

    I try to explain it clearly, only to find ending up with such a long article. :)


Log in to reply
 

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