rb_tree.py

lass Node(object):
    left = None
    right = None
    parent = None
    color = "red"

    def __init__(self, value):
        self.value = value

class Tree(object):
    root = None

    def insert(self, value):
        node = Node(value)
        prev = None
        cur = self.root
        while cur:
            prev = cur
            if node.value < cur.value:
                cur = cur.left
            else:
                cur = cur.right
        node.parent = prev
        if not prev:
            self.root = node
        elif node.value < prev.value:
            prev.left = node
        else:
            prev.right = node

        # 1) node is root:
        #   done
        # 2) node's parent is to the grandparent's left
        #   a) uncle is red:
        #       recolor parent and uncle black, color grandparent red
        #   b) node is to parent's right:
        #       move up to parent, rotate parent left, color parent black, color grandparent red, rotate grandparent right
        #   c) node is to parent's left:
        #       color parent black, color grandparent red, rotate grandparent right
        # 3) node's parent is to the grandparent's right
        #   a) uncle is red:
        #       recolor parent and uncle black, color grandparent red
        #   b) node is to parent's left:
        #       move up to parent, rotate parent right, color parent black, color grandparent red, rotate grandparent left
        #   c) node is to parent's right:
        #       color parent black, color grandparent red, rotate grandparent left
        # we are iterating until we are back at the top of the tree, but why the check on node's parent's color?
        while node != self.root and node.parent.color == "red":
            # how to get the uncle
            if node.parent == node.parent.parent.left:
                uncle = node.parent.parent.right
                # if uncle is red turn it and it's sibling black, turn the grandparent red and move up to it
                if uncle and uncle.color == "red":
                    node.parent.color = "black"
                    uncle.color = "black"
                    node.parent.parent.color = "red"
                else:
                    # if node is to the parent's right rotate left
                    if node == node.parent.right:
                        node = node.parent
                        self.rotate_left(node)
                    node.parent.color = "black"
                    node.parent.parent.color = "red"
                    # rotate the grandparent right
                    self.rotate_right(node.parent.parent)
            else:
                uncle = node.parent.parent.left
                # if uncle is red turn it and it's sibling black, turn the grandparent red and move up to it
                if uncle and uncle.color == "red":
                    node.parent.color = "black"
                    uncle.color = "black"
                    node.parent.parent.color = "red"
                else:
                    # if node is to the parent's left rotate right
                    if node == node.parent.left:
                        node = node.parent
                        self.rotate_right(node)
                    node.parent.color = "black"
                    node.parent.parent.color = "red"
                    # rotate the grandparent left
                    self.rotate_left(node.parnet.parent)
        self.root.color = "black"

    def rotate_right(self, node):
        """Right-rotates node x on tree T.
                n
               / \
              l   r
             / \
            lc rc
        mutates into:
                l
               / \
              lc  n
                 / \
                rc  r
        Used for maintaining tree balance.
        """
        left_child = node.left
        node.left = left_child.right
        left_child.right = node

        if not node.parent:
            self.root = left_child
        elif node == node.parent.right:
            node.parent.right = left_child
        else:
            node.parent.left = left_child

        left_child.parent = node.parent
        if left_child.right:
            left_child.right.parent = node

    def rotate_left(self, node):
        """Left-rotates node x on tree T.
                 n
                / \
               l   r
                  / \
                 lc rc
        mutates into:
                 r
                / \
               n  rc
              / \
             l  lc
        Used for maintaining tree balance.
        """
        right_child = node.right
        right_child.left = node
        node.right = right_child.left

        if not node.parent:
            self.root = right_child
        elif node == node.parent.right:
            node.parent.right = right_child
        else:
            node.parent.left = right_child

        right_child.parent = node.parent
        if right_child.left:
            right_child.left.parent = node

    def in_order(self, node):
        if node.left: self.in_order(node.left)
        print(node.value)
        if node.right: self.in_order(node.right)

    def level_order(self):
        q = []
        q.append(self.root)

        while q:
            v = q.pop(0)
            print(v.value)
            if v.left: q.append(v.left)
            if v.right: q.append(v.right)

t = Tree()
t.insert(7)
t.insert(6)
t.insert(5)
t.insert(4)
t.insert(3)
t.insert(2)
t.insert(1)
print("should be 1 2 3 4 5 6 7")
t.in_order(t.root)
print("should be 6 4 7 2 5 1 3")
t.level_order()