As the note in the problem statement, this problem has a straight-forward O(n)-space solution, which is to generate the inorder traversal results of the bst and then find the two nodes that violate the increasing trend and those are the two that requires to be swapped.
This link utilizes inorder traversal to find those two nodes without keeping all the results. However, inorder traversal implemented recursively will take at least O(logn) space and may even take O(n) space at the worst case. The code is rewritten as follows.
1 class Solution { 2 public: 3 void recoverTree(TreeNode* root) { 4 pre = first = second = NULL; 5 inorder(root); 6 if (first) { 7 int temp = first -> val; 8 first -> val = second -> val; 9 second -> val = temp; 10 } 11 } 12 private: 13 TreeNode *first, *second, *pre; 14 void inorder(TreeNode* root) { 15 if (!root) return; 16 inorder(root -> left); 17 if (pre && pre -> val > root -> val) { 18 if (!first) first = pre; 19 second = root; 20 } 21 pre = root; 22 inorder(root -> right); 23 } 24 };
So to come up with an O(1)-space solution, we indeed have to turn to Morris traversal. This link has a nice explanation of how Morris traversal can be modified to find the two wrong nodes.
1 class Solution { 2 public: 3 void recoverTree(TreeNode* root) { 4 TreeNode* cur = root; 5 TreeNode *first, *second, *pre; 6 first = second = pre = NULL; 7 while (cur) { 8 if (cur -> left) { 9 TreeNode* predecessor = cur -> left; 10 while (predecessor -> right && predecessor -> right != cur) 11 predecessor = predecessor -> right; 12 if (!(predecessor -> right)) { 13 predecessor -> right = cur; 14 cur = cur -> left; 15 } 16 else { 17 predecessor -> right = NULL; 18 if (pre && pre -> val > cur -> val) { 19 if (!first) first = pre; 20 second = cur; 21 } 22 pre = cur; 23 cur = cur -> right; 24 } 25 } 26 else { 27 if (pre && pre -> val > cur -> val) { 28 if (!first) first = pre; 29 second = cur; 30 } 31 pre = cur; 32 cur = cur -> right; 33 } 34 } 35 if (first) swap(first -> val, second -> val); 36 } 37 };