def __balanceAfterRemove__(self, node: TreeNode):
if node.hasParent():
sibling = node.getSibling()
if sibling is not None:
if sibling.isRed():
node.parent.setRedColor()
sibling.setBlackColor()
if node.isLeftChild():
self.turnLeft(node.parent)
else:
self.turnRight(node.parent)
else:
if sibling.leftChild.isBlack() \
and sibling.rightChild.isBlack():
sibling.setRedColor()
if node.parent.isBlack():
self.__balanceAfterRemove__(node.parent)
else:
node.parent.setBlackColor()
return
if sibling.rightChild.isBlack() \
and sibling.leftChild.isRed() \
and node.isLeftChild():
sibling.setRedColor()
sibling.leftChild.setBlackColor()
self.turnRight(sibling)
return
elif sibling.leftChild.isBlack() \
and sibling.rightChild.isRed() \
and node.isRightChild():
sibling.setRedColor()
sibling.rightChild.setBlackColor()
self.turnLeft(sibling)
return
sibling.color = node.parent.color
node.parent.setBlackColor()
if node.isLeftChild():
sibling.rightChild.setBlackColor()
self.turnLeft(node.parent)
else:
sibling.leftChild.setBlackColor()
self.turnRight(node.parent)
def remove(self, node: TreeNode):
if node.hasBothChildren():
newNode = node.rightChild.findMin()
if newNode.isLeftChild():
newNode.parent.leftChild = None
else:
newNode.parent.rightChild = None
if node.isLeftChild():
node.parent.leftChild = newNode
else:
node.parent.rightChild = newNode
newNode.parent = node.parent
newNode.leftChild = node.leftChild
newNode.rightChild = node.rightChild
newNode.color = node.color
elif node.hasAnyChildren():
self.__removeWithSingleChildren__(node)
else:
if node.isLeftChild():
node.parent.leftChild = None
else:
node.parent.rightChild = None
def turnLeft(self, rotateRoot: TreeNode):
newRoot = rotateRoot.rightChild
rotateRoot.rightChild = newRoot.leftChild
if newRoot.leftChild is not None:
newRoot.leftChild.parent = rotateRoot
newRoot.parent = rotateRoot.parent
if rotateRoot.isRoot():
self.root = newRoot
else:
if rotateRoot.isLeftChild():
rotateRoot.parent.leftChild = newRoot
else:
rotateRoot.parent.rightChild = newRoot
newRoot.leftChild = rotateRoot
rotateRoot.parent = newRoot
def __removeWithSingleChildren__(self, node: TreeNode):
successor = node.leftChild if node.hasLeftChild() else node.rightChild
if node.isLeftChild():
node.parent.leftChild = successor
else:
node.parent.rightChild = successor
if node.isRed():
return None
if successor.isRed():
successor.setBlackColor()
else:
self.__balanceAfterRemove__(successor)
def __checkNodeAndParentBranchAndTurn__(self, node: TreeNode):
# Case 2: дядя чёрный (или его нет), папа и дед в разных сторонах
if node.isRightChild() and node.parent.isLeftChild():
self.turnLeft(node.parent)
currentNode = node.leftChild
elif node.isLeftChild() and node.parent.isRightChild():
self.turnRight(node.parent)
currentNode = node.rightChild
else:
currentNode = node
currentNode.parent.setBlackColor()
grandparent = currentNode.getGrandParent()
if grandparent is not None:
grandparent.setRedColor()
# Case 3: дядя чёрный (или его нет), папа и дед в одной стороне
if currentNode.isLeftChild() and currentNode.parent.isLeftChild():
self.turnRight(grandparent)
else:
self.turnLeft(grandparent)