RP2040/P49_RedBlackTree.py at main · ArmstrongSubero/RP2040 · 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
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
# File: main.py
# Author: Armstrong Subero
# Platform: Python (Standard or MicroPython)
# Program: P49_RedBlackTree
# Interpreter: Python 3.x
# Program Version: 1.0
#
# Program Description: This program demonstrates the creation of a Red-Black Tree
#                      with insertion functionality and maintaining the Red-Black properties.
#
# Notes:
# A Red-Black Tree is a balanced binary search tree with an extra bit of storage per node: its color, which can be red or black.

class Node:
    def __init__(self, data):
        self.data = data
        self.parent = None
        self.left = None
        self.right = None
        self.color = 'red'  # New nodes are always red

class RedBlackTree:
    def __init__(self):
        self.TNULL = Node(0)
        self.TNULL.color = 'black'
        self.root = self.TNULL

    def insert(self, key):
        node = Node(key)
        node.left = self.TNULL
        node.right = self.TNULL
        node.color = 'red'

        parent = None
        current = self.root

        while current != self.TNULL:
            parent = current
            if node.data < current.data:
                current = current.left
            else:
                current = current.right

        node.parent = parent

        if parent is None:
            self.root = node
        elif node.data < parent.data:
            parent.left = node
        else:
            parent.right = node

        if node.parent is None:
            node.color = 'black'
            return

        if node.parent.parent is None:
            return

        self._fix_insert(node)

    def _rotate_left(self, x):
        y = x.right
        x.right = y.left
        if y.left != self.TNULL:
            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 _rotate_right(self, x):
        y = x.left
        x.left = y.right
        if y.right != self.TNULL:
            y.right.parent = x

        y.parent = x.parent
        if x.parent is 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 _fix_insert(self, k):
        while k.parent.color == 'red':
            if k.parent == k.parent.parent.right:
                u = k.parent.parent.left
                if u.color == 'red':
                    u.color = 'black'
                    k.parent.color = 'black'
                    k.parent.parent.color = 'red'
                    k = k.parent.parent
                else:
                    if k == k.parent.left:
                        k = k.parent
                        self._rotate_right(k)
                    k.parent.color = 'black'
                    k.parent.parent.color = 'red'
                    self._rotate_left(k.parent.parent)
            else:
                u = k.parent.parent.right
                if u.color == 'red':
                    u.color = 'black'
                    k.parent.color = 'black'
                    k.parent.parent.color = 'red'
                    k = k.parent.parent
                else:
                    if k == k.parent.right:
                        k = k.parent
                        self._rotate_left(k)
                    k.parent.color = 'black'
                    k.parent.parent.color = 'red'
                    self._rotate_right(k.parent.parent)
            if k == self.root:
                break
        self.root.color = 'black'

    def inorder(self):
        self._inorder_helper(self.root)

    def _inorder_helper(self, node):
        if node != self.TNULL:
            self._inorder_helper(node.left)
            print(node.data, end=" ")
            self._inorder_helper(node.right)

# Example usage:

rbt = RedBlackTree()

# Insert some elements
rbt.insert(20)
rbt.insert(15)
rbt.insert(25)
rbt.insert(10)
rbt.insert(5)
rbt.insert(1)

# Perform an inorder traversal
print("Inorder Traversal of Red-Black Tree:")
rbt.inorder()  # Output: 1 5 10 15 20 25