450. Delete Node in a BST

gzc · Oct 3, 2021 · ab73850 · ab73850
1 parent d4a885f
commit ab73850
Showing 1 changed file with 11 additions and 3 deletions.
diff --git a/cpp/441-450/Delete Node in a BST.cpp b/cpp/441-450/Delete Node in a BST.cpp
@@ -7,6 +7,7 @@
  *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
  * };
  */
+
 class Solution {
 public:
     TreeNode* deleteNode(TreeNode* root, int key) {
@@ -20,11 +21,18 @@ class Solution {
                 TreeNode* right = root->right;
                 while (right->left)
                     right = right->left;
-                swap(root->val, right->val);    
+                swap(root->val, right->val);
+                root->right = deleteNode(root->right, key);
+                return root;
             }
         }
-        root->left = deleteNode(root->left, key);
-        root->right = deleteNode(root->right, key);
+
+        if (key < root->val) {
+            root->left = deleteNode(root->left, key);
+        } else {
+            root->right = deleteNode(root->right, key);
+        }
+
         return root;
     }
 };