【一日一题】删除二叉搜索树中的节点

发表于

题目描述
题目来源
给定一个二叉搜索树的根节点 root 和一个值 key，删除二叉搜索树中的 key 对应的节点，并保证二叉搜索树的性质不变。返回二叉搜索树（有可能被更新）的根节点的引用。

一般来说，删除节点可分为两个步骤：

首先找到需要删除的节点；
如果找到了，删除它。
说明： 要求算法时间复杂度为 O(h)，h 为树的高度。

代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    def deleteNode(self, root, key):
        """
        :type root: TreeNode
        :type key: int
        :rtype: TreeNode
        """
        parent = None
        selected = None
        current = root
        while current:
            if current.val == key:
                selected = current
                break
            elif current.val > key:
                if not current.left:
                    break
                parent = current
                current = current.left
            else:
                if not current.right:
                    break
                parent = current
                current = current.right
        if not selected:
            return root
        newSelected = self.rebuildTree(selected.left, selected.right)
        if parent:
            if parent.val > selected.val:
                parent.left = newSelected
            else:
                parent.right = newSelected
        return root if parent else newSelected
    
    def rebuildTree(self, left, right):
        if not left:
            return right
        if not right:
            return left
        current = right
        while current.left:
            current = current.left
        current.left = left
        return right