Tony's Log

Algorithms, Distributed System, Machine Learning

  博客园 :: 首页 :: 博问 :: 闪存 :: 新随笔 :: 联系 :: 订阅 订阅 :: 管理 ::

Not hard to find a solution, but there are several corner cases.

class Solution {
public:
    /**
     * @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) {
        //    Find it
        TreeNode *p = root, *pp = nullptr;
        bool bLeft = false;
        while(p)
        {
            if(p->val == value)
            {
                break;
            }
            pp = p;
            if(p->val < value)
            {
                p = p->right;
            }
            else if(p->val > value)
            {
                p = p->left;
            }
            bLeft = pp->left == p;
        }
        if(!p) return root;

        //    Case 1: no children
        if(!p->left && !p->right)
        {
            if (!pp) return nullptr;
            
            if(bLeft)
                pp->left = nullptr;
            else
                pp->right= nullptr;
            return root;
        }

        //    Compose new left sub-tree
        TreeNode *pl = p->left, *pr = p->right, *pNewSubRoot = nullptr;
        //    Case 2: has both children
        if (pl && pr)
        {
            TreeNode *p0 = pl, *pp0 = pl;
            while(p0->right)
            {
                pp0 = p0;
                p0 = p0->right;
            }
            if(pp0 != p0)
            {
                pp0->right = p0->left;
                p0->left = pl;
            }
            p0->right = pr;

            pNewSubRoot = p0;
        }
        //    Case 3: only 1 child
        else
        {
            TreeNode *pc = pl ? pl : pr;
            if (!pp) return pc;
            if(bLeft)
                pp->left = pc;
            else
                pp->right= pc;
        }

        //    Return
        if(!pp)
        {
            return pNewSubRoot;
        }

        return root;
    }
};
View Code
posted on 2015-10-05 06:35  Tonix  阅读(187)  评论(0编辑  收藏  举报