Fix remove() in binary search tree.

This commit is contained in:
krahets
2023-05-26 20:34:22 +08:00
parent ee716a2c23
commit b39e79be85
9 changed files with 111 additions and 50 deletions

View File

@ -43,7 +43,7 @@ class BinarySearchTree:
def search(self, num: int) -> TreeNode | None:
"""查找节点"""
cur: TreeNode | None = self.__root
cur: TreeNode | None = self.root
# 循环查找,越过叶节点后跳出
while cur is not None:
# 目标节点在 cur 的右子树中
@ -60,11 +60,11 @@ class BinarySearchTree:
def insert(self, num: int) -> None:
"""插入节点"""
# 若树为空,直接提前返回
if self.__root is None:
if self.root is None:
return
# 循环查找,越过叶节点后跳出
cur, pre = self.__root, None
cur, pre = self.root, None
while cur is not None:
# 找到重复节点,直接返回
if cur.val == num:
@ -87,11 +87,11 @@ class BinarySearchTree:
def remove(self, num: int) -> None:
"""删除节点"""
# 若树为空,直接提前返回
if self.__root is None:
if self.root is None:
return
# 循环查找,越过叶节点后跳出
cur, pre = self.__root, None
cur, pre = self.root, None
while cur is not None:
# 找到待删除节点,跳出循环
if cur.val == num:
@ -112,10 +112,14 @@ class BinarySearchTree:
# 当子节点数量 = 0 / 1 时, child = null / 该子节点
child = cur.left or cur.right
# 删除节点 cur
if pre.left == cur:
pre.left = child
if cur != self.root:
if pre.left == cur:
pre.left = child
else:
pre.right = child
else:
pre.right = child
# 若删除节点为根节点,则重新指定根节点
self.__root = cur
# 子节点数量 = 2
else:
# 获取中序遍历中 cur 的下一个节点