Red–black tree impl

gistfile1.py

class Node(object):
    def __init__(self, key, val, color, n):
        self.key = key
        self.val = val
        self.color = color
        self.n = n
        self.left = None
        self.right = None


class RedBlack(object):
    RED = True
    BLACK = False

    def __init__(self):
        self.root = None

    def put(self, key, val):
        self.root = self.put_impl(self.root, key, val)
        self.root.color = self.BLACK

    def put_impl(self, h, key, val):
        if h is None:
            return Node(key, val, self.RED, 1)

        if key < h.key:
            h.left = self.put_impl(h.left, key, val)
        elif key > h.key:
            h.right = self.put_impl(h.right, key, val)
        else:
            h.val = val

        if self.is_red(h.right) and not self.is_red(h.left):
            h = self.rotate_left(h)

        if self.is_red(h.left) and self.is_red(h.left.left):
            h = self.rotate_right(h)

        if self.is_red(h.left) and self.is_red(h.right):
            self.flip_colors(h)

        h.n = self.size(h.left) + self.size(h.right) + 1

        return h

    def is_red(self, x):
        if x is None:
            return False
        return x.color == self.RED

    def rotate_left(self, h):
        x = h.right
        h.right = x.left
        x.left = h
        x.color = x.left.color
        x.left.color = self.RED
        x.n = h.n
        h.n = self.size(h.left) + self.size(h.right) + 1
        return x

    def rotate_right(self, h):
        x = h.left
        h.left = x.right
        x.right = h
        x.color = x.right.color
        x.right.color = self.RED
        x.n = h.n
        h.n = self.size(h.left) + self.size(h.right) + 1
        return x

    def flip_colors(self, h):
        h.color = (not h.color)
        h.left.color = (not h.left.color)
        h.right.color = (not h.right.color)

    def size(self, x):
        if x is None:
            return 0

        return x.n