Leet_Code/0450_delete_node_in_a_bst.py at master · Kcheung42/Leet_Code · GitHub

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
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
#------------------------------------------------------------------------------
# Question:0450_delete_node_in_a_bst.py
#------------------------------------------------------------------------------
# tags:
'''
Given a root node reference of a BST and a key, delete the node with the given
key in the BST. Return the root node reference (possibly updated) of the BST.

Basically, the deletion can be divided into two stages:

Search for a node to remove.
If the node is found, delete the node.
Note: Time complexity should be O(height of tree).

Example:

root = [5,3,6,2,4,null,7]
key = 3

    5
   / \
  3   6
 / \   \
2   4   7

Given key to delete is 3. So we find the node with value 3 and delete it.

One valid answer is [5,4,6,2,null,null,7], shown in the following BST.

    5
   / \
  4   6
 /     \
2       7

Another valid answer is [5,2,6,null,4,null,7].

    5
   / \
  2   6
   \   \
    4   7

         x
       /   \
      x     x
     / \   / \
    x   x x   x

'''

#------------------------------------------------------------------------------
# Solutions
#------------------------------------------------------------------------------
from typing import *
from test_utils.BinaryTree import BinaryTree, TreeNode

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    '''
    Time:
    Space:
    '''
    def findMin(self, root):
        if root.left is None:
            return root.val
        return self.findMin(root.left)

    def deleteNode(self, root: TreeNode, key: int) -> TreeNode:
        if not root:
            return None
        elif key < root.val: root.left = self.deleteNode(root.left, key)
        elif key > root.val: root.right = self.deleteNode(root.right, key)
        else:
            # case 1 no child
            if not(root.right and root.left):
                root = None
            # case 2: One child
            elif root.left and root.right is None:
                temp = root
                root = root.left
                del(temp)
            elif root.right and root.left is None:
                temp = root
                root = root.right
                del(temp)
            # case 3: Two children
            else:
                min_val = self.findMin(root.right)
                root.val = min_val
                root.right = self.deleteNode(root.right, min_val)

        return root


#------------------------------------------------------------------------------
# Tests
#------------------------------------------------------------------------------
import unittest

class TestSolution(unittest.TestCase):
    def inorder(self, root):
        def recur(root, result):
            if root is None:
                return
            recur(root.left, result)
            result.append(root.val)
            recur(root.right, result)

        result = []
        recur(root, result)
        return result

    def test_simple(self):
        tree = [5,3,6,2,4,None,7]
        root = BinaryTree(tree).root
        s = Solution()
        self.assertEqual(self.inorder(s.deleteNode(root, 3)), [2,4,5,6,7])

    def test_simple2(self):
        tree = [1,None,3]
        root = BinaryTree(tree).root
        key = 3
        s = Solution()
        self.assertEqual(self.inorder(s.deleteNode(root, key)), [1])


unittest.main(verbosity=2)