In [1]:

import numpy as np
from bokeh.plotting import figure, show, output_notebook
from bokeh.io import push_notebook
from time import sleep

algorithm¶

In [2]:

class Node:
    def __init__(self, value, left=None, right=None, red=False):
        self.value = value
        self.left = left
        self.right = right
        self.red = red

    @staticmethod
    def red(node):
        return node and node.red

    def rotate_left(self):
        if not Node.red(self.left) and Node.red(self.right):
            child = self.right
            self.right, child.left = child.left, self
            self.red, child.red = True, self.red
            return child
        else:
            return self

    def rotate_right(self):
        if Node.red(self.left) and Node.red(self.left.left):
            child = self.left
            self.left, child.right = child.right, self
            self.red, child.red = True, self.red
            return child
        else:
            return self
    
    def flip_colors(self):
        if Node.red(self.left) and Node.red(self.right):
            self.red, self.left.red, self.right.red = True, False, False

In [3]:

def insert(node, value, root=True):
    if not node:
        return Node(value, red=not root)
    
    # insert value
    if value == node.value:
        pass
    elif value < node.value:
        node.left = insert(node.left, value, root=False)
    else:
        node.right = insert(node.right, value, root=False)
    
    # update tree w.r.t invariants
    node = node.rotate_left()
    node = node.rotate_right()
    node.flip_colors()

    # keep root black
    if root:
        node.red = False

    return node

In [4]:

def draw(node, red, black, xlo=0, xhi=1, x=.5, y=0):
    if node:
        x_, y_ = (xlo + xhi) / 2, y - 1
        (red if node.red else black).append([x, y, x_, y_])
        draw(node.left, red, black, xlo, x_, x_, y_)
        draw(node.right, red, black, x_, xhi, x_, y_)

run¶

In [5]:

# insert few values
root = None
for i in range(5):
    root = insert(root, i)

In [6]:

output_notebook()

# get plot data
red, black = [], []
draw(root, red, black)

# plot
plot = figure(plot_height=400, plot_width=800)
plot.axis.visible = False
plot.grid.visible = False

red_segment = plot.segment(*zip(*red), color='#ff0000')
black_segment = plot.segment(*zip(*black), color='#000000')

handle = show(plot, notebook_handle=True)

Loading BokehJS ...

In [7]:

for _ in range(100):
    for i in np.random.randint(0, 1000, 5):
        root = insert(root, i)

    red, black = [], []
    draw(root, red, black)

    x0, y0, x1, y1 = red and zip(*red) or [tuple()] * 4
    red_segment.data_source.data.update({
        'x0': x0, 'y0': y0, 'x1': x1, 'y1': y1
    })
    x0, y0, x1, y1 = black and zip(*black) or [tuple()] * 4
    black_segment.data_source.data.update({
        'x0': x0, 'y0': y0, 'x1': x1, 'y1': y1
    })

    push_notebook()
    sleep(.1)

In [ ]: