LeetCode – Refresh – Recover Binary Search Tree

 1 /**
 2  * Definition for binary tree
 3  * struct TreeNode {
 4  *     int val;
 5  *     TreeNode *left;
 6  *     TreeNode *right;
 7  *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 8  * };
 9  */
10 class Solution {
11 public:
12     void getTree(vector<int> &result, vector<TreeNode *> &nodes, TreeNode *root) {
13         if (!root) return;
14         getTree(result, nodes, root->left);
15         result.push_back(root->val);
16         nodes.push_back(root);
17         getTree(result, nodes, root->right);
18     }
19     void recoverTree(TreeNode *root) {
20         vector<int> nums;
21         vector<TreeNode *> nodes;
22         getTree(nums, nodes, root);
23         sort(nums.begin(), nums.end());
24         for (int i = 0; i < nums.size(); i++) {
25             nodes[i]->val = nums[i];
26         }
27     }
28 };

 1 /**
 2  * Definition for binary tree
 3  * struct TreeNode {
 4  *     int val;
 5  *     TreeNode *left;
 6  *     TreeNode *right;
 7  *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 8  * };
 9  */
10 class Solution {
11 public:
12     void recoverTree(TreeNode *root) {
13         if (!root) return;
14         TreeNode *n1, *n2, *parent, *prev;
15         bool found = false;
16         while (root) {
17             if (!root->left) {
18                 if (parent && parent->val > root->val) {
19                     if (!found) {
20                         n1 = parent;
21                         found = true;
22                     }
23                     n2 = root;
24                 }
25                 parent = root;
26                 root = root->right;
27             } else {
28                 prev = root->left;
29                 while (prev->right && prev->right != root) prev = prev->right;
30                 if (!prev->right) {
31                     prev->right = root;
32                     root = root->left;
33                 } else {
34                     if (parent && parent->val > root->val) {
35                         if (!found) {
36                             n1 = parent;
37                             found = true;
38                         }
39                         n2 = root;
40                     }
41                     parent = root;
42                     prev->right = NULL;
43                     root = root->right;
44                 }
45             }
46         }
47         if (n1 && n2) {
48             int t = n1->val;
49             n1->val = n2->val;
50             n2->val = t;
51         }
52     }
53 };