lib/rbTree.c

Go to the documentation of this file.

/* rbTree - rbTreeRed-rbTreeBlack Tree - a type of binary tree which 
 * automatically keeps relatively balanced during
 * inserts and deletions.
 *   original author: Shane Saunders
 *   adapted into local conventions: Jim Kent
 */

#include "common.h"
#include "localmem.h"
#include "rbTree.h"

static char const rcsid[] = "$Id: rbTree.c,v 1.10 2007/02/24 15:59:29 kent Exp $";


static struct rbTreeNode *restructure(struct rbTree *t, int tos, 
        struct rbTreeNode *x, struct rbTreeNode *y, struct rbTreeNode *z)
/* General restructuring function - checks for all
 * restructuring cases.  Call when insert has messed up tree.
 * Sadly delete has to do even more work. */
{
struct rbTreeNode *parent, *midNode;

if(y == x->left) 
    {
    if(z == y->left) 
        {  /* in-order:  z, y, x */
        midNode = y;
        y->left = z;
        x->left = y->right;
        y->right = x;
        }
    else 
        {  /* in-order:  y, z, x */
        midNode = z;
        y->right = z->left;
        z->left = y;
        x->left = z->right;
        z->right = x;
        }
    }
else 
    {
    if(z == y->left) 
        {  /* in-order:  x, z, y */
        midNode = z;
        x->right = z->left;
        z->left = x;
        y->left = z->right;
        z->right = y;
        }
    else 
        {  /* in-order:  x, y, z */
        midNode = y;
        x->right = y->left;
        y->left = x;
        y->right = z;
        }
    }
if(tos != 0) 
    {
    parent = t->stack[tos-1];
    if(x == parent->left) 
        parent->left = midNode;
    else 
        parent->right = midNode;
    }
else 
    t->root = midNode;

return midNode;
}

struct rbTree *rbTreeNewDetailed(int (*compare)(void *, void *), struct lm *lm, 
        struct rbTreeNode *stack[128])
/* Allocate rbTree on an existing local memory & stack.  This is for cases
 * where you want a lot of trees, and don't want the overhead for each one. 
 * Note, to clean these up, just do freez(&rbTree) rather than rbFreeTree(&rbTree). */
{
struct rbTree *t;
AllocVar(t);
t->root = NULL;
t->compare = compare;
t->lm = lm;
t->stack = stack;       
t->n = 0;
return t;
}

struct rbTree *rbTreeNew(int (*compare)(void *, void *))
/* rbTreeNew() - Allocates space for a red-black tree and returns a pointer
 * to it.  The function compare compares they keys of two items, and returns a
 * negative, zero, or positive integer depending on whether the first item is
 * less than, equal to, or greater than the second.  */
{
/* The stack keeps us from having to keep explicit
 * parent, grandparent, greatgrandparent variables.
 * It needs to be big enough for the maximum depth
 * of tree.  Since the whole point of rb trees is
 * that they are self-balancing, this is not all
 * that deep, just 2*log2(N).  Therefore a stack of
 * 128 is good for up to 2^64 items in stack, which
 * should keep us for the next couple of decades... */
struct lm *lm = lmInit(0);
struct rbTreeNode **stack = lmAlloc(lm, 128 * sizeof(stack[0]));        
return rbTreeNewDetailed(compare, lm, stack);
}


void rbTreeFree(struct rbTree **pTree)
/* rbTreeFree() - Frees space used by the red-black tree pointed to by t. */
{
struct rbTree *tree = *pTree;
if (tree != NULL)
    {
    lmCleanup(&tree->lm);
    freez(pTree);
    }
}


void *rbTreeAdd(struct rbTree *t, void *item)
/* rbTreeAdd() - Inserts an item into the red-black tree pointed to by t,
 * according the the value its key.  The key of an item in the red-black
 * tree must be unique among items in the tree.  If an item with the same key
 * already exists in the tree, a pointer to that item is returned.  Otherwise,
 * NULL is returned, indicating insertion was successful.
 */
{
struct rbTreeNode *x, *p, *q, *m, **attachX;
int (* compare)(void *, void *);
int cmpResult;
rbTreeColor col;
struct rbTreeNode **stack = NULL;
int tos;

tos = 0;    
if((p = t->root) != NULL) 
    {
    compare = t->compare;
    stack = t->stack;

    /* Repeatedly explore either the left branch or the right branch
     * depending on the value of the key, until an empty branch is chosen.
     */
    for(;;) 
        {
        stack[tos++] = p;
        cmpResult = compare(item, p->item);
        if(cmpResult < 0) 
            {
            p = p->left;
            if(!p) 
                {
                p = stack[--tos];
                attachX = &p->left;
                break;
                }
            }
        else if(cmpResult > 0) 
            {
            p = p->right;
            if(!p) 
                {
                p = stack[--tos];
                attachX = &p->right;
                break;
                }
            }
        else 
            {
            return p->item;
            }
        }
    col = rbTreeRed;
    }
else 
    {
    attachX = &t->root;
    col = rbTreeBlack;
    }

/* Allocate new node and place it in tree. */
if ((x = t->freeList) != NULL)
    t->freeList = x->right;
else
    lmAllocVar(t->lm, x);
x->left = x->right = NULL;
x->item = item;
x->color = col;
*attachX = x;
t->n++;

/* Restructuring or recolouring will be needed if node x and its parent, p,
 * are both red.
 */
if(tos > 0) 
    {
    while(p->color == rbTreeRed) 
        {  /* Double red problem. */

        /* Obtain a pointer to p's parent, m, and sibling, q. */
        m = stack[--tos];
        q = p == m->left ? m->right : m->left;
        
        /* Determine whether restructuring or recolouring is needed. */
        if(!q || q->color == rbTreeBlack) 
            {
            /* Sibling is black.  ==>  Perform restructuring. */
            
            /* Restructure according to the left to right order, of nodes
             * m, p, and x.
             */
            m = restructure(t, tos, m, p, x);
            m->color = rbTreeBlack;
            m->left->color = m->right->color = rbTreeRed;

            /* Restructuring eliminates the double red problem. */
            break;
            }
        /* else just need to flip color */
        
        /* Sibling is also red.  ==>  Perform recolouring. */
        p->color = rbTreeBlack;
        q->color = rbTreeBlack;

        if(tos == 0) break;  /* The root node always remains black. */
            
        m->color = rbTreeRed;

        /* Continue, checking colouring higher up. */
        x = m;
        p = stack[--tos];
        }
    }

return NULL;
}


void *rbTreeFind(struct rbTree *t, void *item)
/* rbTreeFind() - Find an item in the red-black tree with the same key as the
 * item pointed to by `item'.  Returns a pointer to the item found, or NULL
 * if no item was found.
 */
{
struct rbTreeNode *p, *nextP;
int (*compare)(void *, void *) = t->compare;
int cmpResult;
    
/* Repeatedly explore either the left or right branch, depending on the
 * value of the key, until the correct item is found.  */
for (p = t->root; p != NULL; p = nextP)
    {
    cmpResult = compare(item, p->item);
    if(cmpResult < 0) 
        nextP = p->left;
    else if(cmpResult > 0) 
        nextP = p->right;
    else 
        return p->item;
    }
return NULL;
}


void *rbTreeRemove(struct rbTree *t, void *item)
/* rbTreeRemove() - Delete an item in the red-black tree with the same key as
 * the item pointed to by `item'.  Returns a pointer to the deleted item,
 * and NULL if no item was found.
 */
{
struct rbTreeNode *p, *r, *x, *y, *z, *b, *newY;
struct rbTreeNode *m;
rbTreeColor removeCol;
void *returnItem;
int (* compare)(void *, void *);
int cmpResult;
struct rbTreeNode **stack;
int i, tos;


/* Attempt to locate the item to be deleted. */
if((p = t->root)) 
    {
    compare = t->compare;
    stack = t->stack;
    tos = 0;
    
    for(;;) 
        {
        stack[tos++] = p;
        cmpResult = compare(item, p->item);
        if(cmpResult < 0) 
            p = p->left;
        else if(cmpResult > 0) 
            p = p->right;
        else 
            /* Item found. */
            break;
        if(!p) return NULL;
        }
    }
else 
    return NULL;

/* p points to the node to be deleted, and is currently on the top of the
 * stack.
 */
if(!p->left) 
    {
    tos--;  /* Adjust tos to remove p. */
    /* Right child replaces p. */
    if(tos == 0) 
        {
        r = t->root = p->right;
        x = y = NULL;
        }
    else 
        {
        x = stack[--tos];
        if(p == x->left) 
            {
            r = x->left = p->right;
            y = x->right;
            }
        else 
            {
            r = x->right = p->right;
            y = x->left;
            }
        }
    removeCol = p->color;
    }
else if(!p->right) 
    {
    tos--;  /* Adjust tos to remove p. */
    /* Left child replaces p. */
    if(tos == 0) 
        {
        r = t->root = p->left;
        x = y = NULL;
        }
    else 
        {
        x = stack[--tos];
        if(p == x->left) 
            {
            r = x->left = p->left;
            y = x->right;
            }
        else 
            {
            r = x->right = p->left;
            y = x->left;
            }
        }
    removeCol = p->color;
    }
else 
    {
    /* Save p's stack position. */
    i = tos-1;
    
    /* Minimum child, m, in the right subtree replaces p. */
    m = p->right;
    do 
        {
        stack[tos++] = m;
        m = m->left;
        } while(m);
    m = stack[--tos];

    /* Update either the left or right child pointers of p's parent. */
    if(i == 0) 
        {
        t->root = m;
        }
    else 
        {
        x = stack[i-1];  /* p's parent. */
        if(p == x->left) 
            {
            x->left = m;
            }
        else 
            {
            x->right = m;
            }
        }
    
    /* Update the tree part m is removed from, and assign m the child
     * pointers of p (only if m is not the right child of p).
     */
    stack[i] = m;  /* Node m replaces node p on the stack. */
    x = stack[--tos];
    r = m->right;
    if(tos != i) 
        {  /* x is equal to the parent of m. */
        y = x->right;
        x->left = r;
        m->right = p->right;
        }
    else 
        { /* m was the right child of p, and x is equal to m. */
        y = p->left;
        }
    m->left = p->left;

    /* We treat node m as the node which has been removed. */
    removeCol = m->color;
    m->color = p->color;
    }

/* Get return value and reuse the space used by node p. */
returnItem = p->item;
p->right = t->freeList;
t->freeList = p;

t->n--;

/* The pointers x, y, and r point to nodes which may be involved in
 * restructuring and recolouring.
 *  x - the parent of the removed node.
 *  y - the sibling of the removed node.
 *  r - the node which replaced the removed node.
 * From the above code, the next entry off the stack will be the parent of
 * node x.
 */

/* The number of black nodes on paths to all external nodes (NULL child
 * pointers) must remain the same for all paths.  Restructuring or
 * recolouring of nodes may be necessary to enforce this.
 */
if(removeCol == rbTreeBlack) 
    {
    /* Removal of a black node requires some adjustment. */
    
    if(!r || r->color == rbTreeBlack) 
        {
        /* A black node replaced the deleted black node.  Note that
         * external nodes (NULL child pointers) are always black, so
         * if r is NULL it is treated as a black node.
         */

        /* This causes a double-black problem, since node r would need to
         * be coloured double-black in order for the black color on
         * paths through r to remain the same as for other paths.
         */

        /* If r is the root node, the double-black color is not necessary
         * to maintain the color balance.  Otherwise, some adjustment of
         * nearby nodes is needed in order to eliminate the double-black
         * problem.  NOTE:  x points to the parent of r.
         */
        if(x) for(;;) 
            {

            /* There are three adjustment cases:
             *  1.  r's sibling, y, is black and has a red child, z.
             *  2.  r's sibling, y, is black and has two black children.
             *  3.  r's sibling, y, is red.
             */
            if(y->color == rbTreeBlack) 
                {

                /* Note the conditional evaluation for assigning z. */
                if(((z = y->left) && z->color == rbTreeRed) ||
                   ((z = y->right) && z->color == rbTreeRed)) 
                       {                    
                    /* Case 1:  perform a restructuring of nodes x, y, and
                     * z.
                     */
                    
                    b = restructure(t, tos, x, y, z);
                    b->color = x->color;
                    b->left->color = b->right->color = rbTreeBlack;
                    
                    break;
                    }
                else 
                    {
                    /* Case 2:  recolour node y red. */
                    
                    y->color = rbTreeRed;
                    
                    if(x->color == rbTreeRed) 
                        {
                        x->color = rbTreeBlack;
                        break;
                        }
                    /* else */

                    if(tos == 0) break;  /* Root level reached. */
                    /* else */
                    
                    r = x;
                    x = stack[--tos];  /* x <- parent of x. */
                    y = x->left == r ? x->right : x->left;
                    }
                }
            else 
                {
                /* Case 3:  Restructure nodes x, y, and z, where:
                 *  - If node y is the left child of x, then z is the left
                 *    child of y.  Otherwise z is the right child of y.
                 */
                if(x->left == y) 
                    {
                    newY = y->right;
                    z = y->left;
                    }
                else 
                    {
                    newY = y->left;
                    z = y->right;
                    }
                
                restructure(t, tos, x, y, z);
                y->color = rbTreeBlack;
                x->color = rbTreeRed;

                /* Since x has moved down a place in the tree, and y is the
                 * new the parent of x, the stack must be adjusted so that
                 * the parent of x is correctly identified in the next call
                 * to restructure().
                 */
                stack[tos++] = y;

                /* After restructuring, node r has a black sibling, newY,
                 * so either case 1 or case 2 applies.  If case 2 applies
                 * the double-black problem does not reappear.
                 */
                y = newY;
                
                /* Note the conditional evaluation for assigning z. */
                if(((z = y->left) && z->color == rbTreeRed) ||
                   ((z = y->right) && z->color == rbTreeRed)) 
                   {                
                    /* Case 1:  perform a restructuring of nodes x, y, and
                     * z.
                     */
                    
                    b = restructure(t, tos, x, y, z);
                    b->color = rbTreeRed;  /* Since node x was red. */
                    b->left->color = b->right->color = rbTreeBlack;
                    }
                else 
                    {
                    /* Case 2:  recolour node y red. */

                    /* Note that node y is black and node x is red. */
                    
                    y->color = rbTreeRed;
                    x->color = rbTreeBlack;
                    }

                break;
                }
            }
        }
    else 
        {
        /* A red node replaced the deleted black node. */

        /* In this case we can simply color the red node black. */
        r->color = rbTreeBlack;
        }
    }
return returnItem;
}

/* Some variables to help recursively dump tree. */
static int dumpLevel;   /* Indentation level. */
static FILE *dumpFile;  /* Output file */
static void (*dumpIt)(void *item, FILE *f);  /* Item dumper. */

static void rTreeDump(struct rbTreeNode *n)
/* Recursively dump. */
{
if (n == NULL)
    return;
spaceOut(dumpFile, ++dumpLevel * 3);
fprintf(dumpFile, "%c ", (n->color ==  rbTreeRed ? 'r' : 'b'));
dumpIt(n->item, dumpFile);
fprintf(dumpFile, "\n");
rTreeDump(n->left);
rTreeDump(n->right);
--dumpLevel;
}

void rbTreeDump(struct rbTree *tree, FILE *f, 
        void (*dumpItem)(void *item, FILE *f))
/* Dump out rb tree to file, mostly for debugging. */
{
dumpFile = f;
dumpLevel = 0;
dumpIt = dumpItem;
fprintf(f, "rbTreeDump\n");
rTreeDump(tree->root);
}



/* Variables to help recursively traverse tree. */
static void (*doIt)(void *item);
static void *minIt, *maxIt;
static int (*compareIt)(void *, void *);

static void rTreeTraverseRange(struct rbTreeNode *n)
/* Recursively traverse tree in range applying doIt. */
{
if (n != NULL)
   {
   int minCmp = compareIt(n->item, minIt);
   int maxCmp = compareIt(n->item, maxIt);
   if (minCmp >= 0)
       rTreeTraverseRange(n->left);
   if (minCmp >= 0 && maxCmp <= 0)
       doIt(n->item);
   if (maxCmp <= 0)
       rTreeTraverseRange(n->right);
   }
}

static void rTreeTraverse(struct rbTreeNode *n)
/* Recursively traverse full tree applying doIt. */
{
if (n != NULL)
    {
    rTreeTraverse(n->left);
    doIt(n->item);
    rTreeTraverse(n->right);
    }
}

void rbTreeTraverseRange(struct rbTree *tree, void *minItem, void *maxItem,
        void (*doItem)(void *item))
/* Apply doItem function to all items in tree such that
 * minItem <= item <= maxItem */
{
doIt = doItem;
minIt = minItem;
maxIt = maxItem;
compareIt = tree->compare;
rTreeTraverseRange(tree->root);
}

void rbTreeTraverse(struct rbTree *tree, void (*doItem)(void *item))
/* Apply doItem function to all items in tree */
{
doIt = doItem;
compareIt = tree->compare;
rTreeTraverse(tree->root);
}

struct slRef *itList;  /* List of items that rbTreeItemsInRange returns. */

static void addRef(void *item)
/* Add item it itList. */
{
refAdd(&itList, item);
}

struct slRef *rbTreeItemsInRange(struct rbTree *tree, void *minItem, void *maxItem)
/* Return a sorted list of references to items in tree between range.
 * slFreeList this list when done. */
{
itList = NULL;
rbTreeTraverseRange(tree, minItem, maxItem, addRef);
slReverse(&itList);
return itList;
}

struct slRef *rbTreeItems(struct rbTree *tree)
/* Return sorted list of items.  slFreeList this when done.*/
{
itList = NULL;
rbTreeTraverse(tree, addRef);
slReverse(&itList);
return itList;
}

int rbTreeCmpString(void *a, void *b)
/* Set up rbTree so as to work on strings. */
{
return strcmp(a, b);
}

int rbTreeCmpWord(void *a, void *b)     
/* Set up rbTree so as to work on case-insensitive strings. */
{
return differentWord(a,b);
}

---