5---Python/RedBlackTree.py at master · karmusha/5---Python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
class Node:
    def __init__(self, value):
        self.value = value
        self.red = True
        self.left = None
        self.right = None
        self.parent = None


class RedBblackTree:
    def __init__(self):
        self.root = None

    def search(self, value):
        current_node = self.root
        while current_node is not None and value != current_node.value:
            if value < current_node.value:
                current_node = current_node.left
            else:
                current_node = current_node.right
        return current_node

    def insert(self, value):
        node = Node(value)
        if self.root is None:
            node.red = False
            self.root = node
            return 'No root'
        last_node = self.root
        while last_node is not None:
            potential_parent = last_node
            if node.value < last_node.value:
                last_node = last_node.left
            else:
                last_node = last_node.right

        node.parent = potential_parent

        if node.value < node.parent.value:
            node.parent.left = node
        else:
            node.parent.right = node

        node.left = None
        node.right = None
        self.balance_tree(node)

    def balance_tree(self, node):
        try:
            while node.parent.red is True and node is not self.root:
                if node.parent == node.parent.parent.left:
                    uncle = node.parent.parent.right
                    if uncle.red:
                        node.parent.red = False
                        uncle.red = False
                        node.parent.parent.red = True
                        node = node.parent.parent
                    else:
                        if node == node.parent.right:
                            node = node.parent

                else:
                    try:
                        uncle = node.parent.parent.left
                        if uncle.red:
                            node.parent.red = False
                            uncle.red = False
                            node.parent.parent.red = True

                    except AttributeError:
                        print('No uncle')
                        break

            self.root.red = False
        except AttributeError:
            print('Tree done')