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?
O(n) space solution:
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 }
In order to only use O(n) space, we recorde the previous visited node. For the first time that we find a mistak (i.e., curNode.val<=res.pre), the wrong node is the previous one. For the second time that we find a mistake, the wrong node is the current node. We record the two wrong values. Then, we do another travesal, and exchange the wrong values.
NOTE: however, there is a kind of cases is that the two wrong nodes are neighbors in the equivalent list. In such cases, we can only find one mistake, i.e., the first mistake. For example: INPUT: [0,1]. In order to deal with such cases, we also record the current node value when we find the first mistake. At last, if we only find one mistake, then the current node value found in the first mistake is the second wrong value.