Python Red-Black Tree Implementation: Code and Explanation
Here's an implementation of a Red-Black Tree in Python:
class Node:
def __init__(self, key):
self.key = key
self.left = None
self.right = None
self.color = 1
class RedBlackTree:
def __init__(self):
self.null_node = Node(0)
self.null_node.color = 0
self.null_node.left = None
self.null_node.right = None
self.root = self.null_node
def insert(self, key):
node = Node(key)
node.parent = None
node.left = self.null_node
node.right = self.null_node
node.color = 1
y = None
x = self.root
while x != self.null_node:
y = x
if node.key < x.key:
x = x.left
else:
x = x.right
node.parent = y
if y == None:
self.root = node
elif node.key < y.key:
y.left = node
else:
y.right = node
if node.parent == None:
node.color = 0
return
if node.parent.parent == None:
return
self.fix_insert(node)
def fix_insert(self, node):
while node.parent.color == 1:
if node.parent == node.parent.parent.right:
u = node.parent.parent.left
if u.color == 1:
u.color = 0
node.parent.color = 0
node.parent.parent.color = 1
node = node.parent.parent
else:
if node == node.parent.left:
node = node.parent
self.right_rotate(node)
node.parent.color = 0
node.parent.parent.color = 1
self.left_rotate(node.parent.parent)
else:
u = node.parent.parent.right
if u.color == 1:
u.color = 0
node.parent.color = 0
node.parent.parent.color = 1
node = node.parent.parent
else:
if node == node.parent.right:
node = node.parent
self.left_rotate(node)
node.parent.color = 0
node.parent.parent.color = 1
self.right_rotate(node.parent.parent)
if node == self.root:
break
self.root.color = 0
def left_rotate(self, x):
y = x.right
x.right = y.left
if y.left != self.null_node:
y.left.parent = x
y.parent = x.parent
if x.parent == 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, x):
y = x.left
x.left = y.right
if y.right != self.null_node:
y.right.parent = x
y.parent = x.parent
if x.parent == None:
self.root = y
elif x == x.parent.right:
x.parent.right = y
else:
x.parent.left = y
y.right = x
x.parent = y
def delete(self, key):
self.delete_node_helper(self.root, key)
def delete_node_helper(self, node, key):
z = self.null_node
while node != self.null_node:
if node.key == key:
z = node
if node.key <= key:
node = node.right
else:
node = node.left
if z == self.null_node:
print('Couldn't find key in the tree')
return
y = z
y_original_color = y.color
if z.left == self.null_node:
x = z.right
self.transplant(z, z.right)
elif (z.right == self.null_node):
x = z.left
self.transplant(z, z.left)
else:
y = self.minimum(z.right)
y_original_color = y.color
x = y.right
if y.parent == z:
x.parent = y
else:
self.transplant(y, y.right)
y.right = z.right
y.right.parent = y
self.transplant(z, y)
y.left = z.left
y.left.parent = y
y.color = z.color
if y_original_color == 0:
self.fix_delete(x)
def fix_delete(self, x):
while x != self.root and x.color == 0:
if x == x.parent.left:
s = x.parent.right
if s.color == 1:
s.color = 0
x.parent.color = 1
self.left_rotate(x.parent)
s = x.parent.right
if s.left.color == 0 and s.right.color == 0:
s.color = 1
x = x.parent
else:
if s.right.color == 0:
s.left.color = 0
s.color = 1
self.right_rotate(s)
s = x.parent.right
s.color = x.parent.color
x.parent.color = 0
s.right.color = 0
self.left_rotate(x.parent)
x = self.root
else:
s = x.parent.left
if s.color == 1:
s.color = 0
x.parent.color = 1
self.right_rotate(x.parent)
s = x.parent.left
if s.right.color == 0 and s.right.color == 0:
s.color = 1
x = x.parent
else:
if s.left.color == 0:
s.right.color = 0
s.color = 1
self.left_rotate(s)
s = x.parent.left
s.color = x.parent.color
x.parent.color = 0
s.left.color = 0
self.right_rotate(x.parent)
x = self.root
x.color = 0
def minimum(self, node):
while node.left != self.null_node:
node = node.left
return node
def transplant(self, u, v):
if u.parent == None:
self.root = v
elif u == u.parent.left:
u.parent.left = v
else:
u.parent.right = v
v.parent = u.parent
def inorder(self, node):
if node != self.null_node:
self.inorder(node.left)
print(node.key)
self.inorder(node.right)
def search(self, node, key):
if node == self.null_node or key == node.key:
return node
if key < node.key:
return self.search(node.left, key)
return self.search(node.right, key)
tree = RedBlackTree()
tree.insert(7)
tree.insert(3)
tree.insert(18)
tree.insert(10)
tree.insert(22)
tree.insert(8)
tree.insert(11)
tree.insert(26)
tree.insert(2)
tree.insert(6)
tree.insert(13)
tree.inorder(tree.root)
tree.delete(18)
tree.inorder(tree.root)
This code defines two classes: 'Node' and 'RedBlackTree'. The 'Node' class represents a node in the tree and has attributes for the node's key, left and right children, parent, and color. The 'RedBlackTree' class has methods for inserting, deleting, searching, and traversing the tree, as well as methods for fixing the tree after insertions and deletions to maintain the Red-Black properties.
The 'insert' method takes a key as input and inserts a new node with that key into the tree, while maintaining the Red-Black properties. The 'delete' method takes a key as input and deletes the node with that key from the tree, while also maintaining the Red-Black properties. The 'inorder' method traverses the tree in order and prints the keys of each node. The 'search' method takes a key as input and returns the node with that key if it is in the tree, or 'None' if it is not.
To test the code, we create a 'RedBlackTree' object 'tree' and insert some keys. We then use the 'inorder' method to print the keys of the tree in order. We then delete a node with key 18 and print the tree again in order to confirm that the deletion was successful.
原文地址: http://www.cveoy.top/t/topic/lUFF 著作权归作者所有。请勿转载和采集!