Red-Black tree

class Tree: def __init__(self): self.root = None self.leftCount = 0 self.rightCount = 0 def add(self, node, parent): if parent == None: node.color = "black" self.root = node print("Root:", node) elif parent != None: if node.key < parent.key: if parent.left == None: self.leftCount += 1 node.parent = parent print("Node:", node) print("Parent:", parent) parent.left = node self.balance() self.recolor(self.root) else: self.add(node, parent.left) elif node.key > parent.key: if parent.right == None: self.rightCount += 1 node.parent = parent print("Node:", node) print("Parent:", parent) parent.right = node self.balance() self.recolor(self.root) else: self.add(node, parent.right) def recolor(self, node): self.root.color = "black" if node.parent != None: if node.parent.color == "black": node.color = "red" if node.parent.color == "red": node.color = "black" if node.left != None: self.recolor(node.left) if node.right != None: self.recolor(node.right) def balance(self): if self.leftCount + self.rightCount > 1: if self.rightCount + 1 < self.leftCount: self.rightRotate() elif self.rightCount > self.leftCount + 1: self.leftRotate() def leftRotate(self): new_root = self.root.right left_sub = self.root.right.left new_root.left = self.root new_root.left.right = left_sub new_root.left.parent = new_root new_root.parent = None self.root = new_root self.leftCount += 1 self.rightCount -= 1 def rightRotate(self): new_root = self.root.left right_sub = self.root.left.right new_root.right = self.root new_root.right.left = right_sub new_root.right.parent = new_root new_root.parent = None self.root = new_root self.leftCount -= 1 self.rightCount += 1 def maximum(self, node): if node.right == None: print("Maximum:", node) print("Color:", node.color) print("-----------------") elif node.right != None: self.maximum(node.right) def predecessor(self, s, parent): try: if s == parent.key: print("Predecessor ({0}) == {1}".format(s, parent.parent.key)) print("Predecessor's color:", parent.parent.color) if s < parent.key: self.predecessor(s, parent.left) elif s > parent.key: self.predecessor(s, parent.right) except AttributeError: print("Or it is a vertex and we can't find predecessor") print("or we haven't got this element.") class Node: def __init__(self, key): self.key = key self.parent = None self.left = None self.right = None self.color = None def __str__(self): return str(self.key) def main(): print("Red-white trees.") print("") lenKeys = int(input("Please enter count of vetices:")) lenKeysSearch = int(input("Please enter count of attempts to find vertices:")) print("------------------------") theTree = Tree() for i in range(lenKeys): key = int(input("Enter key:")) node = Node(key) theTree.add(node, theTree.root) theTree.maximum(theTree.root) for i in range(lenKeysSearch): key = int(input("Enter key for searching (predecessor):")) theTree.predecessor(key, theTree.root) if __name__ == "__main__": main() input()