Online Judge Solutions

Saturday, December 6, 2014

LintCode: Remove Node in Binary Search Tree

Given a root of Binary Search Tree with unique value for each node.  Remove the node with given value. If there is no such a node with given value in the binary search tree, do nothing. You should keep the tree still a binary search tree after removal.
Example
Given binary search tree:
          5
       /    \
    3          6
 /    \
2       4
Remove 3, you can either return:
          5
       /    \
    2          6
      \
         4
or :
          5
       /    \
    4          6
 /   
2
 
/**
 * Definition of TreeNode:
 * class TreeNode {
 * public:
 *     int val;
 *     TreeNode *left, *right;
 *     TreeNode(int val) {
 *         this->val = val;
 *         this->left = this->right = NULL;
 *     }
 * }
 */
class Solution {
public:
   TreeNode* removeNode(TreeNode* root, TreeNode *parent, int value) {
         
        if (!root) return NULL;
         
        if (root->val == value) {
            if (parent == NULL) {
                TreeNode parent(0);
                parent.left = root;
                removeNode(root, &parent, value);
                return parent.left;
            }
            else if (!root->left|| !root->right) {
                TreeNode *desc= root->left?root->left: root->right;
                if (root == parent->left)
                    parent->left = desc;
                else 
                    parent->right = desc;
            }
            else {
                TreeNode *right = root->right;
                while(right->left) right = right->left;
                
                root->val = right->val;
                removeNode(root->right, root, root->val);
            }
        }
        else if (root->val <  value) 
            removeNode(root->right, root, value);   
        else 
            removeNode(root->left, root, value);
            
        return root;
    };
    
    /**
     * @param root: The root of the binary search tree.
     * @param value: Remove the node with given value.
     * @return: The root of the binary search tree after removal.
     */
    TreeNode* removeNode(TreeNode* root, int value) {
        return removeNode(root, NULL, value);
    }
};

No comments:

Post a Comment