Left-learnig Red Black Tree

http://hillbig.cocolog-nifty.com/do/2008/11/post-b6ee.html

ざっと見て、かなり簡単に書けそうだったのでPythonで書いてみました。
ほとんどそのまんまでサクッと動いた。これは簡単。
要素の検索、追加、削除ができます。

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

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

    def __init__(self):
        self.root = None

    def get(self, key):
        return get(self.root, key)

    def put(self, key, value):
        self.root = insert(self.root, key, value)

    def delete_max(self):
        self.root = delete_max(self.root)
        self.root.color = BinarySearchTree.BLACK

    def delete_min(self):
        self.root = delete_min(self.root)
        self.root.color = BinarySearchTree.BLACK

    def delete(self, key):
        self.root = delete(self.root, key)

def get(node, key):
    x = node 
    while x:
        c = cmp(key, x.key)
        if c == 0:
            return x.value
        elif c < 0:
            x = x.left
        else:
            x = x.right
    return None

def is_red(node):
    if node: 
        return node.color == BinarySearchTree.RED
    else: 
        return False

def rotate_left(node):
    h = node
    x = h.right
    h.right = x.left
    x.left = h
    x.color = x.left.color
    x.left.color = BinarySearchTree.RED
    return x

def rotate_right(node):
    h = node
    x = h.left
    h.left = x.right
    x.right = h
    x.color = x.right.color
    x.right.color = BinarySearchTree.RED
    return x

def color_flip(node):
    h = node
    h.color = not h.color
    h.left.color = not h.left.color
    h.right.color = not h.right.color
    return h 

def insert(node, key, value):
    h = node
    if not h:
        return BinarySearchTree.Node(key, value, BinarySearchTree.RED)

    c = cmp(key, h.key)
    if c == 0:
        h.value = value
    elif (c < 0):
        h.left = insert(h.left, key, value)
    else:
        h.right = insert(h.right, key, value)

    if is_red(h.right):
        h = rotate_left(h)
    if is_red(h.left) and is_red(h.left.left):
        h = rotate_right(h)
    if is_red(h.left) and is_red(h.right):
        color_flip(h)

    return h

def fix_up(node):
    h = node
    if is_red(h.right):
        h = rotate_left(h)
    if is_red(h.left) and is_red(h.left.left):
        h = rotate_right(h)
    if is_red(h.left) and is_red(h.right):
        color_flip(h)
    return h

def move_red_right(node):
    h = node
    color_flip(h)
    if is_red(h.left.left):
        h = rotate_right(h)
        color_flip(h)
    return h

def move_red_left(node):
    h = node
    color_flip(h)
    if is_red(h.right.left):
        h.right = rotate_right(h.right)
        h = rotate_left(h)
        color_flip(h)
    return h

def delete_max(node):
    h = node
    if is_red(h.left):
        h = rotate_right(h)

    if not h.right:
        return None

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

    h.left = delete_max(h.left)

    return fix_up(h)

def delete_min(node):
    h = node
    if not h.left:
        return None

    if not is_red(h.left) and not is_red(h.left.left):
        h = move_red_left(h)

    h.left = delete_min(h.left)

    return fix_up(h)

def delete(node, key):
    h = node
    c = cmp(key, h.key)
    if c < 0:
        if not is_red(h.left) and not is_red(h.left.left):
            h = move_red_left(h)
        h.left = delete(h.left, key)
    else:
        if is_red(h.left):
            h = lean_right(h)
        if c == 0 and not h.right:
            return None

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

        if c == 0:
            h.key = min(h.right)
            h.value = get(h.right, h.key)
            h.right = delete_min(h.right)
        else:
            h.right = delete(h.right, key)

    return fix_up(h)



if __name__ == '__main__':
    bst = BinarySearchTree()
    bst.put('apple', 100)
    bst.put('banana', 200)
    bst.put('cherry', 300)
    bst.put('orange', 400)
    bst.put('melon', 500)
    bst.put('mango', 600)

    print 'apple', bst.get('apple')
    print 'banana', bst.get('banana')
    print 'cherry', bst.get('cherry')
    print 'orange', bst.get('orange')
    print 'melon', bst.get('melon')
    print 'mango', bst.get('mango')

    bst.delete('cherry')
    print 'apple', bst.get('apple')
    print 'banana', bst.get('banana')
    print 'cherry', bst.get('cherry')
    print 'orange', bst.get('orange')
    print 'melon', bst.get('melon')
    print 'mango', bst.get('mango')

定番のようなアルゴリズムでもまだまだ工夫の余地があるものなんですね。
複雑さに面食らうようなものでも、実はじっくり考え直すと新しい発見があるのかもしれない。