rbt.cpp

多功能的红黑树和压缩算法

// RBT.cpp: implementation of the CRBT class.
//
//////////////////////////////////////////////////////////////////////

// reentrant red-black tree
#include "public.h"
#include "rbt.h"

//////////////////////////////////////////////////////////////////////
// Construction/Destruction
//////////////////////////////////////////////////////////////////////


//##ModelId=444B252100B0


// all leafs are sentinels
#define SENTINEL &rbt->sentinel

CRBT::CRBT()
{
	m_pTreeHandle = NULL;
	memset(m_errMsg,0,sizeof(m_errMsg));
}

CRBT::~CRBT()
{
	
}


bool CRBT::rbtNew(int(*rbtCompare)(void *a, void *b)) 
{
	if (m_pTreeHandle != NULL)
	{
		return false;
	}
    RbtType *rbt;
    
    if ((m_pTreeHandle = (RbtType *)malloc(sizeof(RbtType))) == NULL) 
	{
        return NULL;
    }
	rbt = (RbtType *)m_pTreeHandle;
    rbt->compare = rbtCompare;
    rbt->root = SENTINEL;
    rbt->sentinel.left = SENTINEL;
    rbt->sentinel.right = SENTINEL;
    rbt->sentinel.parent = NULL;
    rbt->sentinel.color = BLACK;
    rbt->sentinel.key = NULL;
    rbt->sentinel.val = NULL;

    return rbt;
}

void CRBT::deleteTree(NodeType *p) 
{
    RbtType *rbt = (RbtType *)m_pTreeHandle;
    // erase nodes depth-first
    if (p == SENTINEL) return;
    deleteTree(p->left);
    deleteTree(p->right);
    free(p);
}

void CRBT::rbtDelete() 
{
	if (m_pTreeHandle == NULL)
	{
		return;
	}
    RbtType *rbt = (RbtType *)m_pTreeHandle;

    deleteTree(rbt->root);
    free(rbt);
	m_pTreeHandle = NULL;
	return;
}

void CRBT::rotateLeft(NodeType *x) 
{

    // rotate node x to left
	RbtType *rbt = (RbtType *)m_pTreeHandle;
    NodeType *y = x->right;

    // establish x->right link
    x->right = y->left;
    if (y->left != SENTINEL) y->left->parent = x;

    // establish y->parent link
    if (y != SENTINEL) y->parent = x->parent;
    if (x->parent) {
        if (x == x->parent->left)
            x->parent->left = y;
        else
            x->parent->right = y;
    } else {
        rbt->root = y;
    }

    // link x and y
    y->left = x;
    if (x != SENTINEL) x->parent = y;
}

void CRBT::rotateRight(NodeType *x) 
{
	RbtType *rbt = (RbtType *)m_pTreeHandle;
    // rotate node x to right

    NodeType *y = x->left;

    // establish x->left link
    x->left = y->right;
    if (y->right != SENTINEL) y->right->parent = x;

    // establish y->parent link
    if (y != SENTINEL) y->parent = x->parent;
    if (x->parent) {
        if (x == x->parent->right)
            x->parent->right = y;
        else
            x->parent->left = y;
    } else {
        rbt->root = y;
    }

    // link x and y
    y->right = x;
    if (x != SENTINEL) x->parent = y;
}

void CRBT::insertFixup(NodeType *x) 
{
   
   // maintain red-black tree balance after inserting node x
   RbtType *rbt = (RbtType *)m_pTreeHandle;
   // check red-black properties
   while (x != rbt->root && x->parent->color == RED) 
   {
      // we have a violation
      if (x->parent == x->parent->parent->left) 
	  {
         NodeType *y = x->parent->parent->right;
         if (y->color == RED) 
		 {
            
            // uncle is RED
            x->parent->color = BLACK;
            y->color = BLACK;
            x->parent->parent->color = RED;
            x = x->parent->parent;
         } 
		 else 
		 {
            
            // uncle is BLACK
            if (x == x->parent->right) {
               // make x a left child
               x = x->parent;
               rotateLeft(x);
            }
            
            // recolor and rotate
            x->parent->color = BLACK;
            x->parent->parent->color = RED;
            rotateRight(x->parent->parent);
         }
      } 
	  else 
	  {
         
         // mirror image of above code
         NodeType *y = x->parent->parent->left;
         if (y->color == RED) 
		 {
            
            // uncle is RED
            x->parent->color = BLACK;
            y->color = BLACK;
            x->parent->parent->color = RED;
            x = x->parent->parent;
         } 
		 else 
		 {
            
            // uncle is BLACK
            if (x == x->parent->left) 
			{
               x = x->parent;
               rotateRight(x);
            }
            x->parent->color = BLACK;
            x->parent->parent->color = RED;
            rotateLeft(x->parent->parent);
         }
      }
   }
   rbt->root->color = BLACK;
}

RbtStatus CRBT::rbtInsert(void *key, void *val) 
{
    NodeType *current, *parent, *x;
    RbtType *rbt = (RbtType *)m_pTreeHandle;

    // allocate node for data and insert in tree

    // find future parent
    current = rbt->root;
    parent = 0;
    while (current != SENTINEL) 
	{
        int rc = rbt->compare(key, current->key);
        if (rc == 0)
            return RBT_STATUS_DUPLICATE_KEY;
        parent = current;
        current = (rc < 0) ? current->left : current->right;
    }

    // setup new node
    if ((x = (struct NodeTag *)malloc (sizeof(*x))) == 0)
	{
        return RBT_STATUS_MEM_EXHAUSTED;
	}
    x->parent = parent;
    x->left = SENTINEL;
    x->right = SENTINEL;
    x->color = RED;
    x->key = key;
    x->val = val;

    // insert node in tree
    if(parent) 
	{
        if(rbt->compare(key, parent->key) < 0)
		{
            parent->left = x;
		}
        else
		{
            parent->right = x;
		}
    } 
	else 
	{
        rbt->root = x;
    }

    insertFixup(x);

    return RBT_STATUS_OK;
}

void CRBT::deleteFixup(NodeType *x) 
{
	RbtType *rbt = (RbtType *)m_pTreeHandle;
    // maintain red-black tree balance after deleting node x

    while (x != rbt->root && x->color == BLACK) 
	{
        if (x == x->parent->left) 
		{
            NodeType *w = x->parent->right;
            if (w->color == RED) 
			{
                w->color = BLACK;
                x->parent->color = RED;
                rotateLeft(x->parent);
                w = x->parent->right;
            }
            if (w->left->color == BLACK && w->right->color == BLACK) 
			{
                w->color = RED;
                x = x->parent;
            } 
			else 
			{
                if (w->right->color == BLACK) 
				{
                    w->left->color = BLACK;
                    w->color = RED;
                    rotateRight(w);
                    w = x->parent->right;
                }
                w->color = x->parent->color;
                x->parent->color = BLACK;
                w->right->color = BLACK;
                rotateLeft(x->parent);
                x = rbt->root;
            }
        } 
		else 
		{
            NodeType *w = x->parent->left;
            if (w->color == RED) 
			{
                w->color = BLACK;
                x->parent->color = RED;
                rotateRight(x->parent);
                w = x->parent->left;
            }
            if (w->right->color == BLACK && w->left->color == BLACK) 
			{
                w->color = RED;
                x = x->parent;
            } 
			else 
			{
                if (w->left->color == BLACK) 
				{
                    w->right->color = BLACK;
                    w->color = RED;
                    rotateLeft(w);
                    w = x->parent->left;
                }
                w->color = x->parent->color;
                x->parent->color = BLACK;
                w->left->color = BLACK;
                rotateRight(x->parent);
                x = rbt->root;
            }
        }
    }
    x->color = BLACK;
}

RbtStatus CRBT::rbtErase(RbtIterator i) 
{
    NodeType *x, *y;
    RbtType *rbt = (RbtType *)m_pTreeHandle;
    NodeType *z = (NodeType *)i;

    if (z->left == SENTINEL || z->right == SENTINEL) 
	{
        // y has a SENTINEL node as a child
        y = z;
    } 
	else 
	{
        // find tree successor with a SENTINEL node as a child
        y = z->right;
        while (y->left != SENTINEL) 
		{
			y = y->left;
		}
    }

    // x is y's only child
    if (y->left != SENTINEL)
	{
        x = y->left;
	}
    else
	{
        x = y->right;
	}

    // remove y from the parent chain
    x->parent = y->parent;
    if (y->parent)
	{
        if (y == y->parent->left)
		{
            y->parent->left = x;
		}
        else
		{
            y->parent->right = x;
		}
	}
    else
	{
        rbt->root = x;
	}
    if (y != z) 
	{
        z->key = y->key;
        z->val = y->val;
    }


    if (y->color == BLACK)
	{
        deleteFixup(x);
	}
    free (y);
	y= NULL;
    return RBT_STATUS_OK;
}

RbtIterator CRBT::rbtNext(RbtIterator it) 
{
    RbtType *rbt = (RbtType *)m_pTreeHandle;
    NodeType *i = (NodeType *)it;

    if (i->right != SENTINEL) 
	{
        // go right 1, then left to the end
        for (i = i->right; i->left != SENTINEL; i = i->left);
    } 
	else 
	{
        // while you're the right child, chain up parent link
        NodeType *p = i->parent;
        while (p && i == p->right) 
		{
            i = p;
            p = p->parent;
        }

        // return the "inorder" node
        i = p;
    }
    return i != SENTINEL ? i : NULL;
}


// 函数名: rbtPrev
// 描述  : 取it前一个索引单元
// 返回  : RbtIterator 
// 参数  : RbtHandle h
// 参数  : RbtIterator it
RbtIterator CRBT::rbtPrev(RbtIterator it) 
{
   RbtType *rbt = (RbtType *)m_pTreeHandle;
   NodeType *i = (NodeType *) it;
   
   if (i->left != SENTINEL) 
   {
      // go right 1, then left to the end
      for (i = i->left; i->right != SENTINEL; i = i->right);
   } 
   else 
   {
      // while you're the left child, chain up parent link
      NodeType *p = i->parent;
      while (p && i == p->left) 
      {
         i = p;
         p = p->parent;
      }
      
      // return the "inorder" node
      i = p;
   }
   return i==SENTINEL?NULL:i;
}




RbtIterator CRBT::rbtBegin() 
{
    RbtType *rbt = (RbtType *)m_pTreeHandle;

    // return pointer to first value
    NodeType *i;
    for (i = rbt->root; i->left != SENTINEL; i = i->left);
    return i != SENTINEL ? i : NULL;
}

RbtIterator CRBT::rbtEnd() 
{
   // return pointer to one past last value
   return NULL;
}

void CRBT::rbtKeyValue(RbtIterator it, void **key, void **val) 
{
    NodeType *i = (NodeType *)it;

    *key = i->key;
    *val = i->val;
}


RbtIterator CRBT::rbtFind(void *key) 
{
    RbtType *rbt = (RbtType *)m_pTreeHandle;

    NodeType *current;
    current = rbt->root;
    while(current != SENTINEL) 
	{
        int rc = rbt->compare(key, current->key);
        if (rc == 0) 
		{
			return current;
		}
        current = (rc < 0) ? current->left : current->right;
    }
    return NULL;
}