Red-Black Trees

October 2, 2009

We represent a red-black tree as a five-slot vector containing the color, key, value, left child and right child. Various access functions are shown below; it would be better, for purposes of speed, to make these macros or to write them inline, but we’ll leave them as-is for clarity:

(define (tree c k v l r) (vector c k v l r))
(define empty (tree 'black 'nil 'nil 'nil 'nil))
(define (empty? t) (eqv? t empty))
(define (color t) (vector-ref t 0))
(define (key t) (vector-ref t 1))
(define (value t) (vector-ref t 2))
(define (lkid t) (vector-ref t 3))
(define (rkid t) (vector-ref t 4))
(define (red? c) (eqv? c 'red))
(define (black? c) (eqv? c 'black))

The lookup function is simple recursion over the tree; it returns #f if the key is not in the tree or a pair containing the key and value if it is:

(define (lookup lt? t k)
  (cond ((empty? t) #f)
        ((lt? k (key t)) (lookup lt? (lkid t) k))
        ((lt? (key t) k) (lookup lt? (rkid t) k))
        (else (cons (key t) (value t)))))

The insert function uses an auxiliary function ins that uses the same four predicates as the lookup function; it inserts a new node if it reaches nil, replaces the current value if it finds the key already in the tree, and otherwise calls balancing operations that recursively insert the new node in the proper child. The body of the function forces the root to be black:

(define (insert lt? t k v)
  (define (ins t)
    (cond ((empty? t) (tree 'red k v empty empty))
          ((lt? k (key t))
            (balance (color t) (key t) (value t) (ins (lkid t)) (rkid t)))
          ((lt? (key t) k)
            (balance (color t) (key t) (value t) (lkid t) (ins (rkid t))))
          (else (tree (color t) k v (lkid t) (rkid t)))))
  (let ((z (ins t)))
    (tree 'black (key z) (value z) (lkid z) (rkid z))))

The balance function looks for black-red-red paths in four possible configurations (red left child has red left grand-child or red right grand-child, or red right child has red left grand-child or red right grand-child) and performs the appropriate rotations in the tree; if you examine each case separately, it is far simpler than it looks:

(define (balance c k v l r)
  (cond ((and (black? c) (red? (color l)) (red? (color (lkid l))))
          (tree 'red (key l) (value l)
            (tree 'black (key (lkid l)) (value (lkid l))
              (lkid (lkid l)) (rkid (lkid l)))
            (tree 'black k v (rkid l) r)))
        ((and (black? c) (red? (color l)) (red? (color (rkid l))))
          (tree 'red (key (rkid l)) (value (rkid l))
            (tree 'black (key l) (value l) (lkid l) (lkid (rkid l)))
            (tree 'black k v (rkid (rkid l)) r)))
        ((and (black? c) (red? (color r)) (red? (color (lkid r))))
          (tree 'red (key (lkid r)) (value (lkid r))
            (tree 'black k v l (lkid (lkid r)))
            (tree 'black (key r) (value r) (rkid (lkid r)) (rkid r))))
        ((and (black? c) (red? (color r)) (red? (color (rkid r))))
          (tree 'red (key r) (value r)
            (tree 'black k v l (lkid r))
            (tree 'black (key (rkid r)) (value (rkid r))
              (lkid (rkid r)) (rkid (rkid r)))))
        (else (tree c k v l r))))

Enlist performs an infix traversal of the tree:

(define (enlist t)
  (let enlist ((t t) (xs '()))
    (cond ((empty? t) xs)
          ((and (empty? (lkid t)) (empty? (rkid t)))
            (cons (cons (key t) (value t)) xs))
          (else (enlist (lkid t)
                        (cons (cons (key t) (value t))
                              (enlist (rkid t) xs)))))))

A sample dialog is shown below:

> (define t
    (insert <
      (insert <
        (insert <
          (insert <
            (insert <
              empty
              2 "b")
            5 "e")
          3 "c")
        4 "d")
      1 "a"))
> (lookup < t 7)
#f
> (lookup < t 4)
(4 . "d")
> (enlist t)
((1 . "a") (2 . "b") (3 . "c") (4 . "d") (5 . "e"))

You can run this program at http://programmingpraxis.codepad.org/UA6qxeO7.

Pages: 1 2

7 Responses to “Red-Black Trees”

  1. Matt Might said

    I needed red-black trees for some static analysis work, but I wanted to be able to add tags to internal nodes.

    Here’s the implementation I came up with for Scala based on the Okasaki book:

    http://matt.might.net/articles/implementation-of-immutable-purely-functional-okasaki-red-black-tree-maps-in-scala/

  2. Vikas Tandi said
    #include <stdlib.h>
    
    typedef struct RedBlackTreeNode
    {
    	int							key;
    	int							color;
    	struct	RedBlackTreeNode*	left;
    	struct	RedBlackTreeNode*	right;
    }rbtree;
    
    static rbtree* RedBlackTree_insert_imp(rbtree *p, int key, int direction);
    static rbtree* rightRotation(rbtree *p);
    static rbtree* leftRotation(rbtree *p);
    static rbtree* create_node(int key);
    
    enum {BLACK = 0, RED};
    enum {LEFT = 0, RIGHT};
    
    /* insert function */
    rbtree* RedBlackTree_insert(rbtree *p, int key)
    {
    	rbtree *root;
    
    	root = RedBlackTree_insert_imp(p, key, LEFT);
    
    	/* the color of root should always be black */
    	root->color = BLACK;
    	return root;
    }
    
    /* lookup function */
    rbtree* RedBlackTree_find(rbtree *p, int key)
    {
    	while(p)
    	{
    		if(p->key == key)
    			return p;
    		else if(key < p->key)
    			p = p->left;
    		else
    			p = p->right;
    	}
    	return NULL;
    }
    
    void RedBlackTree_free(rbtree *p)
    {
    	if(p == NULL)
    		return;
    
    	RedBlackTree_free(p->left);
    	RedBlackTree_free(p->right);
    	free(p);
    }
    
    static rbtree* RedBlackTree_insert_imp(rbtree *p, int key, int direction)
    {
    	if(p == NULL)
    		return create_node(key);
    
    	/* while moving from top to bottom, check for 4-nodes and split them*/
    	if(p->left != NULL && p->right != NULL)
    	{
    		if((p->left->color == RED) && (p->right->color == RED))
    		{
    			/* split the 4-node into two 2-nodes by changing the color of left and right
    				child from red to black. Also, convert the parent of 4-node from n-node to (n+1)-node */
    			p->color = RED;
    			p->left->color = BLACK;
    			p->right->color = BLACK;
    		}
    	}
    
    	if(key < p->key)
    	{
    		/* move to left subtree */
    		p->left = RedBlackTree_insert_imp(p->left, key, LEFT);
    
    		/* three nodes are connected by red links. the nodes orientation is n1--right--p--left-n3.
    			rotate p right */
    		if(p->color == RED			&&
    			p->left->color == RED	&&
    			direction == RIGHT)
    		{
    			p = rightRotation(p);
    		}
    
    		/* three nodes are connected by red links. the nodes orientation is n1--left--p--left-n3.
    			rotate p right */         
    		if(p->left != NULL			&&
    			p->left->left != NULL	&&
    			p->left->color == RED	&&
    			p->left->left->color == RED)
    		{
    			p = rightRotation(p);
    
    			/* change the color of p to black and right child to red */
    			p->color = BLACK;
    			p->right->color = RED;
    		}
    	}
    	else
    	{
    		/* move to right subtree */
    		p->right = RedBlackTree_insert_imp(p->right, key, RIGHT);
    
    		/* three nodes are connected by red links. the nodes orientation is n1--left--p--right-n3.
    			rotate p left */
    		if(p->color == RED			&&
    			p->right->color == RED	&&
    			direction == LEFT)
    		{
    			p = leftRotation(p);
    		}
    
    		/* three nodes are connected by red links. the nodes orientation is n1--right--p--right-n3.
    			rotate p left */
    		if(p->right != NULL			&&
    			p->right->right != NULL	&&
    			p->right->color == RED	&&
    			p->right->right->color == RED)
    		{
    			p = leftRotation(p);
    
    			/* change the color of p to black and right child to red */
    			p->color = BLACK;
    			p->left->color = RED;
    		}
    	}
    
    	return p;
    }
    
    static rbtree* create_node(int key)
    {
    	rbtree *p;
    
    	p = (rbtree*)malloc(sizeof(*p));
    
    	if(p == NULL)
    		return NULL;
    
    	p->color = RED;
    	p->key = key;
    	p->left = p->right = NULL;
    	return p;
    }
    
    static rbtree* rightRotation(rbtree *p)
    {
    	rbtree *ptr = p->left;
    
    	p->left = ptr->right;
    	ptr->right = p;
    
    	return ptr;
    }
    
    static rbtree* leftRotation(rbtree *p)
    {
    	rbtree *ptr = p->right;
    
    	p->right = ptr->left;
    	ptr->left = p;
    
    	return ptr;
    }
    
  3. ;; runs in Chicken Scheme, however portable to most Scheme implementations that have pattern matching (most do)

    (module red-black-tree
    (key value color left right make-node search insert traverse-inorder
    traverse-preorder traverse-postorder leaf)

    (import chicken scheme matchable data-structures)
    (require-library matchable utf8-srfi-13)

    (define-syntax define*
    (syntax-rules ()
    ((_ name body …) (define name (match-lambda* body …)))))

    ;; _ :: [string] -> string
    (define ++ string-append)

    ;; _ :: node a b -> maybe a
    (define (key n) (vector-ref n 0))

    ;; _ :: node a b -> maybe b
    (define (value n) (vector-ref n 1))

    ;; _ :: node a b -> maybe symbol
    (define (color n) (vector-ref n 2))

    ;; _ :: node a b -> maybe node a b
    (define (left n) (vector-ref n 3))
    (define (right n) (vector-ref n 4))

    ;; _ :: a -> b -> symbol -> node a b -> node a b -> node a b
    (define (make-node k v c lc rc) (vector k v c lc rc))

    ;; _ :: node a b
    (define (leaf) (make-node ‘null ‘null ‘B ‘null ‘null))
    (define root leaf)

    ;; _ :: a -> b -> symbol -> node a b
    (define (make-leaves key val color) (make-node key val color (leaf) (leaf)))

    ;; _ :: node a b -> boolean
    (define (leaf? n) (equal? n (leaf)))

    ;; _ :: node a b -> a
    (define* search
    ((T K) (search T K node a b
    (define (recolor-parent n)
    (make-node (key n) (value n) ‘B (left n) (right n)))

    ;; _ :: node a b -> a -> b -> function-symbol -> node a
    (define* insert
    ((T K V) (insert T K V node a b
    (define* balance
    ;; 1) red left child has red left grandchild
    (#(K V ‘B #(K* V* ‘R #(K** V** ‘R L** R**) R*) R)
    (make-node K* V* ‘R (make-node K** V** ‘B L** R**) (make-node K V ‘B R* R)))
    ;; 2) red left child has red right grandchild
    (#(K V ‘B #(K* V* ‘R L* #(K** V** ‘R L** R**)) R)
    (make-node K** V** ‘R (make-node K* V* ‘B L* L**) (make-node K V ‘B R** R)))
    ;; 3) red right child has red left grandchild
    (#(K V ‘B L #(K* V* ‘R #(K** V** ‘R L** R**) R*))
    (make-node K** V** ‘R (make-node K V ‘B L L**) (make-node K* V* ‘B R** R*)))
    ;; 4) red right child has red right grandchild
    (#(K V ‘B L #(K* V* ‘R L* #(K** V** ‘R L** R**)))
    (make-node K* V* ‘R (make-node K V ‘B L L*) (make-node K** V** ‘B L** R**)))
    ((T) T))

    ;; _ :: node a b -> maybe IO ()
    (define (traverse-inorder t)
    (if (leaf? t) “”
    (let ((key* (->string (key t))) (val (->string (value t)))
    (color* (->string (color t))))
    (traverse-inorder (left t))
    (display (++ “(” key* “,” val “,” color* “) “))
    (traverse-inorder (right t)))))
    (define (traverse-preorder t)
    (if (leaf? t) “”
    (let ((key* (->string (key t))) (val (->string (value t)))
    (color* (->string (color t))))
    (display (++ “(” key* “,” val “,” color* “) “))
    (traverse-preorder (left t))
    (traverse-preorder (right t)))))
    (define (traverse-postorder t)
    (if (leaf? t) “”
    (let ((key* (->string (key t))) (val (->string (value t)))
    (color* (->string (color t))))
    (traverse-preorder (left t))
    (traverse-preorder (right t))
    (display (++ “(” key* “,” val “,” color* “) “)))))

    ;; example:
    ;;(define t root)
    ;;(define t (insert t 2 “b” <))
    ;;(display (++ (traverse-inorder t) "\n"))
    ;;(define t (insert t 5 "e" <))
    ;;(display (++ (traverse-inorder t) "\n"))
    ;;(define t (insert t 3 "c" <))
    ;;(display (++ (traverse-inorder t) "\n"))
    ;;(define t (insert t 4 "d" <))
    ;;(display (++ (traverse-inorder t) "\n"))
    ;;(define t (insert t 1 "a" <))
    ;;(display (++ (traverse-inorder t) "\n"))

    )

  4. sorry, it didn’t post it correctly due to formatting errors, here is a link:

    http://beyert.dyndns.org/files/src/red-black-tree.scm

  5. Here is a more permanent link, please delete my other posts:

    http://codepad.org/YiMdrEem

  6. Sorry about all of the extra posts, this one has a saved link on codepad, and the proper semi-permanent link on my webpage…

    http://codepad.org/HETfVVJz

    http://beyert.dyndns.org/src/red-black-tree.scm

  7. […] functional data structures. The example that really interested me was Okasaki’s functional Red-Black tree. The insertion and balance routines were so short and elegant that I felt I had to be able to […]

Leave a comment