297 lines
11 KiB
Python
297 lines
11 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
Реализация красно-чёрного дерева (Red-Black Tree) на Python.
|
||
|
||
Красно-чёрное дерево — это самобалансирующееся бинарное дерево поиска,
|
||
которое гарантирует логарифмическое время выполнения операций вставки,
|
||
удаления и поиска.
|
||
|
||
Свойства красно-чёрного дерева:
|
||
1. Каждый узел либо красный, либо чёрный.
|
||
2. Корень дерева всегда чёрный.
|
||
3. Все листья (NIL) чёрные.
|
||
4. Если узел красный, то оба его потомка чёрные.
|
||
5. Для каждого узла все пути от него до листьев содержат одинаковое
|
||
количество чёрных узлов.
|
||
"""
|
||
|
||
class Node:
|
||
def __init__(self, key, color="red"):
|
||
self.key = key
|
||
self.left = None
|
||
self.right = None
|
||
self.parent = None
|
||
self.color = color # "red" или "black"
|
||
|
||
|
||
class RedBlackTree:
|
||
def __init__(self):
|
||
self.NIL = Node(None, "black") # Листовые узлы
|
||
self.root = self.NIL
|
||
|
||
def insert(self, key):
|
||
"""Вставка нового узла в дерево."""
|
||
new_node = Node(key)
|
||
new_node.left = self.NIL
|
||
new_node.right = self.NIL
|
||
|
||
parent = None
|
||
current = self.root
|
||
|
||
# Поиск места для вставки
|
||
while current != self.NIL:
|
||
parent = current
|
||
if new_node.key < current.key:
|
||
current = current.left
|
||
else:
|
||
current = current.right
|
||
|
||
new_node.parent = parent
|
||
if parent is None:
|
||
self.root = new_node
|
||
elif new_node.key < parent.key:
|
||
parent.left = new_node
|
||
else:
|
||
parent.right = new_node
|
||
|
||
# Исправление свойств дерева
|
||
self._fix_insert(new_node)
|
||
|
||
def _fix_insert(self, node):
|
||
"""Исправление свойств дерева после вставки."""
|
||
while node != self.root and node.parent.color == "red":
|
||
if node.parent == node.parent.parent.left:
|
||
uncle = node.parent.parent.right
|
||
if uncle.color == "red":
|
||
# Случай 1: Дядя красный
|
||
node.parent.color = "black"
|
||
uncle.color = "black"
|
||
node.parent.parent.color = "red"
|
||
node = node.parent.parent
|
||
else:
|
||
if node == node.parent.right:
|
||
# Случай 2: Дядя чёрный, узел — правый потомок
|
||
node = node.parent
|
||
self._left_rotate(node)
|
||
# Случай 3: Дядя чёрный, узел — левый потомок
|
||
node.parent.color = "black"
|
||
node.parent.parent.color = "red"
|
||
self._right_rotate(node.parent.parent)
|
||
else:
|
||
uncle = node.parent.parent.left
|
||
if uncle.color == "red":
|
||
# Случай 1: Дядя красный
|
||
node.parent.color = "black"
|
||
uncle.color = "black"
|
||
node.parent.parent.color = "red"
|
||
node = node.parent.parent
|
||
else:
|
||
if node == node.parent.left:
|
||
# Случай 2: Дядя чёрный, узел — левый потомок
|
||
node = node.parent
|
||
self._right_rotate(node)
|
||
# Случай 3: Дядя чёрный, узел — правый потомок
|
||
node.parent.color = "black"
|
||
node.parent.parent.color = "red"
|
||
self._left_rotate(node.parent.parent)
|
||
|
||
self.root.color = "black"
|
||
|
||
def _left_rotate(self, x):
|
||
"""Левый поворот вокруг узла x."""
|
||
y = x.right
|
||
x.right = y.left
|
||
if y.left != self.NIL:
|
||
y.left.parent = x
|
||
y.parent = x.parent
|
||
if x.parent is None:
|
||
self.root = y
|
||
elif x == x.parent.left:
|
||
x.parent.left = y
|
||
else:
|
||
x.parent.right = y
|
||
y.left = x
|
||
x.parent = y
|
||
|
||
def _right_rotate(self, y):
|
||
"""Правый поворот вокруг узла y."""
|
||
x = y.left
|
||
y.left = x.right
|
||
if x.right != self.NIL:
|
||
x.right.parent = y
|
||
x.parent = y.parent
|
||
if y.parent is None:
|
||
self.root = x
|
||
elif y == y.parent.left:
|
||
y.parent.left = x
|
||
else:
|
||
y.parent.right = x
|
||
x.right = y
|
||
y.parent = x
|
||
|
||
def search(self, key):
|
||
"""Поиск узла по ключу."""
|
||
current = self.root
|
||
while current != self.NIL and key != current.key:
|
||
if key < current.key:
|
||
current = current.left
|
||
else:
|
||
current = current.right
|
||
return current if current != self.NIL else None
|
||
|
||
def delete(self, key):
|
||
"""Удаление узла по ключу."""
|
||
node = self.search(key)
|
||
if node is None:
|
||
return
|
||
|
||
original_color = node.color
|
||
if node.left == self.NIL:
|
||
x = node.right
|
||
self._transplant(node, node.right)
|
||
elif node.right == self.NIL:
|
||
x = node.left
|
||
self._transplant(node, node.left)
|
||
else:
|
||
successor = self._minimum(node.right)
|
||
original_color = successor.color
|
||
x = successor.right
|
||
if successor.parent == node:
|
||
x.parent = successor
|
||
else:
|
||
self._transplant(successor, successor.right)
|
||
successor.right = node.right
|
||
successor.right.parent = successor
|
||
|
||
self._transplant(node, successor)
|
||
successor.left = node.left
|
||
successor.left.parent = successor
|
||
successor.color = node.color
|
||
|
||
if original_color == "black":
|
||
self._fix_delete(x)
|
||
|
||
def _transplant(self, u, v):
|
||
"""Замена поддерева u поддеревом v."""
|
||
if u.parent is None:
|
||
self.root = v
|
||
elif u == u.parent.left:
|
||
u.parent.left = v
|
||
else:
|
||
u.parent.right = v
|
||
v.parent = u.parent
|
||
|
||
def _minimum(self, node):
|
||
"""Поиск узла с минимальным ключом в поддереве."""
|
||
while node.left != self.NIL:
|
||
node = node.left
|
||
return node
|
||
|
||
def _fix_delete(self, x):
|
||
"""Исправление свойств дерева после удаления."""
|
||
while x != self.root and x.color == "black":
|
||
if x == x.parent.left:
|
||
sibling = x.parent.right
|
||
if sibling.color == "red":
|
||
sibling.color = "black"
|
||
x.parent.color = "red"
|
||
self._left_rotate(x.parent)
|
||
sibling = x.parent.right
|
||
|
||
if sibling.left.color == "black" and sibling.right.color == "black":
|
||
sibling.color = "red"
|
||
x = x.parent
|
||
else:
|
||
if sibling.right.color == "black":
|
||
sibling.left.color = "black"
|
||
sibling.color = "red"
|
||
self._right_rotate(sibling)
|
||
sibling = x.parent.right
|
||
|
||
sibling.color = x.parent.color
|
||
x.parent.color = "black"
|
||
sibling.right.color = "black"
|
||
self._left_rotate(x.parent)
|
||
x = self.root
|
||
else:
|
||
sibling = x.parent.left
|
||
if sibling.color == "red":
|
||
sibling.color = "black"
|
||
x.parent.color = "red"
|
||
self._right_rotate(x.parent)
|
||
sibling = x.parent.left
|
||
|
||
if sibling.right.color == "black" and sibling.left.color == "black":
|
||
sibling.color = "red"
|
||
x = x.parent
|
||
else:
|
||
if sibling.left.color == "black":
|
||
sibling.right.color = "black"
|
||
sibling.color = "red"
|
||
self._left_rotate(sibling)
|
||
sibling = x.parent.left
|
||
|
||
sibling.color = x.parent.color
|
||
x.parent.color = "black"
|
||
sibling.left.color = "black"
|
||
self._right_rotate(x.parent)
|
||
x = self.root
|
||
|
||
x.color = "black"
|
||
|
||
def inorder_traversal(self, node=None):
|
||
"""Обход дерева в порядке возрастания."""
|
||
if node is None:
|
||
node = self.root
|
||
result = []
|
||
if node != self.NIL:
|
||
result += self.inorder_traversal(node.left)
|
||
result.append(node.key)
|
||
result += self.inorder_traversal(node.right)
|
||
return result
|
||
|
||
def print_tree(self, node=None, indent="", last=True):
|
||
"""Вывод дерева в консоль (для визуализации)."""
|
||
if node is None:
|
||
node = self.root
|
||
if node != self.NIL:
|
||
print(indent, end="")
|
||
if last:
|
||
print("R----", end="")
|
||
indent += " "
|
||
else:
|
||
print("L----", end="")
|
||
indent += "| "
|
||
|
||
color = node.color
|
||
print(f"{node.key}({color})")
|
||
self.print_tree(node.left, indent, False)
|
||
self.print_tree(node.right, indent, True)
|
||
|
||
|
||
# Пример использования
|
||
if __name__ == "__main__":
|
||
rbt = RedBlackTree()
|
||
keys = [7, 3, 18, 10, 22, 8, 11, 26]
|
||
for key in keys:
|
||
rbt.insert(key)
|
||
|
||
print("Обход дерева в порядке возрастания:")
|
||
print(rbt.inorder_traversal())
|
||
|
||
print("\nСтруктура дерева:")
|
||
rbt.print_tree()
|
||
|
||
print("\nПоиск узла с ключом 10:")
|
||
node = rbt.search(10)
|
||
print(f"Найден узел: {node.key} ({node.color})")
|
||
|
||
print("\nУдаление узла с ключом 18:")
|
||
rbt.delete(18)
|
||
print("Обход дерева после удаления:")
|
||
print(rbt.inorder_traversal())
|
||
|
||
print("\nСтруктура дерева после удаления:")
|
||
rbt.print_tree()
|