Leetcode-Recover BST

  1 /**
  2  * Definition for binary tree
  3  * public class TreeNode {
  4  *     int val;
  5  *     TreeNode left;
  6  *     TreeNode right;
  7  *     TreeNode(int x) { val = x; }
  8  * }
  9  */
 10 
 11 class SearchRes{
 12     int foundMis;
 13     int firstMisVal;
 14     int secondMisVal;
 15     int firstMisAfter;
 16     int pre;
 17     boolean start;
 18     
 19     SearchRes(){
 20         foundMis = 0;
 21         firstMisVal = -1;
 22         secondMisVal = -1;
 23         firstMisAfter = -1;
 24         pre = -1;
 25         start = false;
 26     }
 27 }
 28 
 29 public class Solution {
 30     public void recoverTree(TreeNode root) {
 31         SearchRes res = new SearchRes();
 32         searchMistakes(root,res);
 33         if (res.foundMis==1)
 34             res.secondMisVal = res.firstMisAfter;
 35         fixMistakes(root,res);
 36     }
 37     
 38     public void fixMistakes(TreeNode curNode, SearchRes res){
 39         if (curNode.left==null&&curNode.right==null){
 40             if (curNode.val==res.firstMisVal)
 41                 curNode.val = res.secondMisVal;
 42             else if (curNode.val==res.secondMisVal)
 43                 curNode.val = res.firstMisVal;
 44             return;
 45         }
 46         
 47         if (curNode.left!=null)
 48             fixMistakes(curNode.left,res);
 49         
 50         if (curNode.val==res.firstMisVal)
 51             curNode.val = res.secondMisVal;
 52         else if (curNode.val==res.secondMisVal)
 53             curNode.val = res.firstMisVal;
 54         
 55         if (curNode.right!=null) 
 56             fixMistakes(curNode.right,res);
 57         return;
 58     }
 59     
 60     public void searchMistakes(TreeNode curNode, SearchRes res){
 61         if (curNode.left==null&&curNode.right==null){
 62             if (!res.start){
 63                 res.start = true;
 64                 res.pre = curNode.val;
 65                 return;
 66             } else {
 67                 //if current value less than pre value, then find a mistake.
 68                 //If this is the first time we find a mistake, then the wrong value is the pre one;
 69                 //otherwise, it is the current one.
 70                 if (curNode.val<=res.pre){
 71                     if (res.foundMis==0){
 72                         res.firstMisVal = res.pre;
 73                         res.firstMisAfter = curNode.val;
 74                     } else 
 75                         res.secondMisVal = curNode.val;
 76                     res.foundMis++;
 77                 }
 78                 res.pre = curNode.val;
 79                 return;
 80             }
 81         }
 82         
 83         
 84         if (curNode.left!=null)
 85             searchMistakes(curNode.left,res);
 86         
 87         if (res.foundMis==2)
 88             return;
 89         
 90         //check current node.
 91         if (!res.start){
 92             res.start = true;
 93             res.pre = curNode.val;
 94         } else {
 95             //if current value less than pre value, then find a mistake.
 96             //If this is the first time we find a mistake, then the wrong value is the pre one;
 97             //otherwise, it is the current one.
 98             if (curNode.val<=res.pre){
 99                 if (res.foundMis==0){
100                     res.firstMisVal = res.pre;
101                     res.firstMisAfter = curNode.val;
102                 } else 
103                     res.secondMisVal = curNode.val;
104                 res.foundMis++;
105             }
106             res.pre = curNode.val;
107         }
108         
109         if (res.foundMis==2)
110             return;
111             
112         if (curNode.right!=null)
113             searchMistakes(curNode.right,res);
114         
115         return;
116     }
117 }