Ir para o conteúdo

Red-black Tree

Aqui está codigo que implementa uma red-black tree. O tempo de inserção e remoção de elementos é de \(O(lg(N))\), onde \(N\) é o número de elementos na árvore. A red-black tree assegura isto mantendo algumas propriedades, através de rotação de árvores, etc.

O código depende do xalloc. De qualquer maneira, a única coisa que faz é verificar o resultado das funções de alocação de memória.

O código foi directamente "traduzido" do pseudo-código do livro Introduction to Algorithms, 2nd Edition, por Thomas H. Cormen, Charles E. Leiserson, Ronald L. Rivest e Clifford Stein, página ~240.

rbtree.h:

/*
 *  GPL
 *
 *  Written by Diogo Sousa aka orium
 *  Copyright (C) 2008 Diogo Sousa
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#ifndef RBTREE_H_
#define RBTREE_H_

enum color
{
    BLACK=1,
    RED
};

struct rbtree_node
{
    void *data;

    enum color color;

    struct rbtree_node *parent;

    struct rbtree_node *left;
    struct rbtree_node *right;
};

struct rbtree
{
    struct rbtree_node *root;

    struct rbtree_node *nil;
};

extern void rbtree_init(struct rbtree **);
extern int rbtree_insert(struct rbtree *, void *, 
             int (*)(const void *, const void *));
extern int rbtree_delete(struct rbtree *, void *,
             int (*)(const void *, const void *));
extern void rbtree_walk(struct rbtree *,
            void (*)(void *));
extern void rbtree_free(struct rbtree *,
            void (*)(void *));
extern void *rbtree_search(struct rbtree *, void *,
               int (*)(const void *, const void *));
extern int rbtree_height(struct rbtree *);

#endif

rbtree.c:

/*
 *  GPL
 *
 *  Written by Diogo Sousa aka orium
 *  Copyright (C) 2008 Diogo Sousa
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <stdlib.h>
#include <assert.h>
#include "rbtree.h"
#include "xalloc.h"

/*
 * This algorithms where directly translated from Introduction to Algorithms,
 * 2nd Edition, by Thomas H. Cormen, Charles E. Leiserson, Ronald L. Rivest
 * and Clifford Stein, page ~240.
 * 
 */

static struct rbtree_node *alloc_rbtree_node(void);
static void rotate_left(struct rbtree *, struct rbtree_node *);
static void rotate_right(struct rbtree *, struct rbtree_node *);
static void insert_fixup(struct rbtree *, struct rbtree_node *);
static struct rbtree_node *rbtree_minimum(struct rbtree *, 
                      struct rbtree_node *);
static struct rbtree_node *rbtree_successor(struct rbtree *, 
                        struct rbtree_node *);
static void delete_fixup(struct rbtree *, struct rbtree_node *);
static void rbtree_walk_recursive(struct rbtree *, struct rbtree_node *,
                  void (*)(void *));
static void rbtree_free_recursive(struct rbtree *, struct rbtree_node *, 
                  void (*)(void *));
static struct rbtree_node *search_node(struct rbtree *, void *,
                       int (*)(const void *, const void *));
static int rbtree_height_recursive(struct rbtree *, struct rbtree_node *,
                   int);

void
rbtree_init(struct rbtree **rbtree)
{
    *rbtree=xmalloc(sizeof(**rbtree));

    (*rbtree)->nil=alloc_rbtree_node();
    (*rbtree)->nil->color=BLACK;
    (*rbtree)->nil->parent=(*rbtree)->nil;
    (*rbtree)->nil->left=(*rbtree)->nil;
    (*rbtree)->nil->right=(*rbtree)->nil;
    (*rbtree)->nil->data=NULL;

    (*rbtree)->root=(*rbtree)->nil;
}

static struct rbtree_node *
alloc_rbtree_node(void)
{
    struct rbtree_node *r;

    r=xmalloc(sizeof(*r));

    return r;
}

/*
 * Before rotate_left(X):
 *
 *         X
 *       /    
 *      A     Y
 *          /    
 *         B     K
 *
 * After rotate_left(X):
 * 
 *         Y
 *       /    
 *      X     K
 *    /    
 *   A     B
 *
 */
static void
rotate_left(struct rbtree *rbtree, struct rbtree_node *x)
{
    struct rbtree_node *y;

    assert(x->right != rbtree->nil);

    y=x->right;
    x->right=y->left;

    y->left->parent=x;
    y->parent=x->parent;

    if (x->parent == rbtree->nil)
        rbtree->root=y;
    else
        if (x == x->parent->left)
            x->parent->left=y;
        else
            x->parent->right=y;

    y->left=x;
    x->parent=y;
}

/*
 * Before rotate_right(Y):
 * 
 *         Y
 *       /    
 *      X     K
 *    /    
 *   A     B
 *
 * After rotate_right(Y):
 *
 *         X
 *       /    
 *      A     Y
 *          /    
 *         B     K
 *
 */
static void
rotate_right(struct rbtree *rbtree, struct rbtree_node *x)
{
    struct rbtree_node *y;

    assert(x->left != rbtree->nil);

    y=x->left;
    x->left=y->right;

    y->right->parent=x;
    y->parent=x->parent;

    if (x->parent == rbtree->nil)
        rbtree->root=y;
    else
        if (x == x->parent->left)
            x->parent->left=y;
        else
            x->parent->right=y;

    y->right=x;
    x->parent=y;
}

int 
rbtree_insert(struct rbtree *rbtree,
          void *data, 
          int (*cmp)(const void *, const void *))
{
    struct rbtree_node *node;
    struct rbtree_node *y;
    struct rbtree_node *x;
    int d;

    if (data == NULL) /* This may be used as a special value */
        return -1;

    node=alloc_rbtree_node();
    node->data=data;

    node->color=RED;
    node->left=rbtree->nil;
    node->right=rbtree->nil;

    x=rbtree->root;
    y=rbtree->nil;

    while (x != rbtree->nil)
    {
        y=x;

        d=cmp(node->data,x->data);

        if (!d)
            return -1;

        if (d < 0)
            x=x->left;
        else
            x=x->right;
    }

    node->parent=y;

    if (y == rbtree->nil)
        rbtree->root=node;
    else
    {
        /* Here d == cmp(node->data,y->data) */

        if (d < 0)
            y->left=node;
        else
            y->right=node;
    }

    insert_fixup(rbtree,node);

    return 0;
}

static void
insert_fixup(struct rbtree *rbtree, struct rbtree_node *node)
{
    struct rbtree_node *y;

    while (node->parent->color == RED)
    {
        if (node->parent == node->parent->parent->left)
        {
            y=node->parent->parent->right;

            if (y->color == RED)
            {
                node->parent->color=BLACK;
                y->color=BLACK;
                node->parent->parent->color=RED;
                node=node->parent->parent;
            }
            else 
            {
                if (node == node->parent->right)
                {
                    node=node->parent;
                    rotate_left(rbtree,node);
                }

                node->parent->color=BLACK;
                node->parent->parent->color=RED;

                rotate_right(rbtree,node->parent->parent);
            }
        }
        else
        {
            y=node->parent->parent->left;

            if (y->color == RED)
            {
                node->parent->color=BLACK;
                y->color=BLACK;
                node->parent->parent->color=RED;
                node=node->parent->parent;
            }
            else 
            {
                if (node == node->parent->left)
                {
                    node=node->parent;
                    rotate_right(rbtree,node);
                }

                node->parent->color=BLACK;
                node->parent->parent->color=RED;

                rotate_left(rbtree,node->parent->parent);
            }

        }
    }

    rbtree->root->color=BLACK;  
}

static struct rbtree_node *
rbtree_minimum(struct rbtree *rbtree,
           struct rbtree_node *node)
{
    while (node->left != rbtree->nil)
        node=node->left;

    return node;
}

static struct rbtree_node *
rbtree_successor(struct rbtree *rbtree,
         struct rbtree_node *node)
{
    struct rbtree_node *y;

    if (node->right != rbtree->nil)
        return rbtree_minimum(rbtree,node->right);

    y=node->parent;

    while (y != rbtree->nil && node == y->right)
    {
        node=y;
        y=y->parent;
    }

    return y;
}

int
rbtree_delete(struct rbtree *rbtree, void *needle, 
          int (*cmp)(const void *, const void *))
{
    struct rbtree_node *node;
    struct rbtree_node *y;
    struct rbtree_node *x;

    node=search_node(rbtree,needle,cmp);

    if (node == rbtree->nil)
        return -1;

    if (node->left == rbtree->nil 
        || node->right == rbtree->nil)
        y=node;
    else
        y=rbtree_successor(rbtree,node);

    if (y->left != rbtree->nil)
        x=y->left;
    else
        x=y->right;

    x->parent=y->parent;

    if (y->parent == rbtree->nil)
        rbtree->root=x;
    else
    {
        if (y == y->parent->left)
            y->parent->left=x;
        else
            y->parent->right=x;
    }

    if (y != node)
        node->data=y->data;

    if (y->color == BLACK)
        delete_fixup(rbtree,x);

    free(y);

    return 0;
}

static void
delete_fixup(struct rbtree *rbtree, struct rbtree_node *node)
{
    struct rbtree_node *w;

    while (node != rbtree->root && node->color == BLACK)
    {
        if (node == node->parent->left)
        {
            w=node->parent->right;

            if (w->color == RED)
            {
                w->color=BLACK;
                node->parent->color=RED;
                rotate_left(rbtree,node->parent);
                w=node->parent->right;
            }

            if (w->left->color == BLACK 
                && w->right->color == BLACK)
            {
                w->color=RED;
                node=node->parent;
            }
            else
            {
                if (w->right->color == BLACK)
                {
                    w->left->color=BLACK;
                    w->color=RED;
                    rotate_right(rbtree,w);
                    w=node->parent->right;
                }

                w->color=node->parent->color;
                node->parent->color=BLACK;
                w->right->color=BLACK;
                rotate_left(rbtree,node->parent);
                node=rbtree->root;
            }
        }
        else
        {
            w=node->parent->left;

            if (w->color == RED)
            {
                w->color=BLACK;
                node->parent->color=RED;
                rotate_right(rbtree,node->parent);
                w=node->parent->left;
            }

            if (w->right->color == BLACK 
                && w->left->color == BLACK)
            {
                w->color=RED;
                node=node->parent;
            }
            else
            {
                if (w->left->color == BLACK)
                {
                    w->right->color=BLACK;
                    w->color=RED;
                    rotate_left(rbtree,w);
                    w=node->parent->left;
                }

                w->color=node->parent->color;
                node->parent->color=BLACK;
                w->right->color=BLACK;
                rotate_right(rbtree,node->parent);
                node=rbtree->root;
            }

        }
    }

    node->color=BLACK;
}

static void 
rbtree_walk_recursive(struct rbtree *rbtree, struct rbtree_node *node,
              void (*action)(void *))
{
    if (node == rbtree->nil)
        return;

    rbtree_walk_recursive(rbtree,node->left,action);
    action(node->data);
    rbtree_walk_recursive(rbtree,node->right,action);
}

void 
rbtree_walk(struct rbtree *rbtree,
        void (*action)(void *))
{
    rbtree_walk_recursive(rbtree,rbtree->root,action);
}

static void 
rbtree_free_recursive(struct rbtree *rbtree, struct rbtree_node *node,
              void (*action)(void *))
{
    if (node == rbtree->nil)
        return;

    if (action != NULL)
        action(node->data);

    rbtree_free_recursive(rbtree,node->left,action);
    rbtree_free_recursive(rbtree,node->right,action);

    free(node);
}

void 
rbtree_free(struct rbtree *rbtree, void (*action)(void *))
{
    rbtree_free_recursive(rbtree,rbtree->root,action);
    free(rbtree->nil);
    free(rbtree);
}

static struct rbtree_node *
search_node(struct rbtree *rbtree, void *needle,
        int (*cmp)(const void *, const void *))
{
    struct rbtree_node *node;
    int d;

    node=rbtree->root;

    while (node != rbtree->nil)
    {
        d=cmp(needle,node->data);

        if (!d)
            break;

        if (d < 0)
            node=node->left;
        else
            node=node->right;
    }

    return node;
}

void *
rbtree_search(struct rbtree *rbtree, void *needle,
          int (*cmp)(const void *, const void *))
{
    struct rbtree_node *node;

    node=search_node(rbtree,needle,cmp);

    return (node == rbtree->nil) ? NULL : node->data;
}

static int
rbtree_height_recursive(struct rbtree *rbtree, struct rbtree_node *node,
            int depth)
{
    int depth_right;
    int depth_left;

    if (node == rbtree->nil)
        return depth-1;

    depth_left=rbtree_height_recursive(rbtree,node->left,depth+1);
    depth_right=rbtree_height_recursive(rbtree,node->right,depth+1);

    return (depth_left > depth_right) ? depth_left : depth_right;
}

int
rbtree_height(struct rbtree *rbtree)
{
    return rbtree_height_recursive(rbtree,rbtree->root,0);
}

Exemplo:

/*
 *  GPL
 *
 *  Written by Diogo Sousa aka orium
 *  Copyright (C) 2008 Diogo Sousa
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <stdio.h>
#include <assert.h>

#include "rbtree.h"

struct integer
{
    int n;
    char *english;
};

int
my_cmp(const void *p1, const void *p2)
{
    return ((struct integer *)p1)->n-((struct integer *)p2)->n;
}

void
my_twalk_action(void *data)
{
    printf("%d %sn",((struct integer *)data)->n,
           ((struct integer *)data)->english);
}

int 
main(void)
{
    struct rbtree *rbtree;
    struct integer i1={1,"one"};
    struct integer i2={2,"two"};
    struct integer i3={3,"three"};  
    struct integer i4={4,"four"};
    struct integer i5={5,"five"};
    struct integer lost={2,NULL};
    struct integer *needle;

    rbtree_init(&rbtree);

    rbtree_insert(rbtree,&i1,my_cmp);
    rbtree_insert(rbtree,&i2,my_cmp);
    rbtree_insert(rbtree,&i3,my_cmp);
    rbtree_insert(rbtree,&i4,my_cmp);
    rbtree_insert(rbtree,&i5,my_cmp);

    rbtree_walk(rbtree,my_twalk_action);

#if 1
    rbtree_delete(rbtree,&lost,my_cmp);
    putchar('n');
    rbtree_walk(rbtree,my_twalk_action);
    putchar('n');
#endif

    needle=rbtree_search(rbtree,&lost,my_cmp);

    if (needle == NULL)
    {
        printf("Not foundn");
        rbtree_free(rbtree,NULL);
        return 0;
    }

    printf("%sn",needle->english);

    printf("height: %dn",rbtree_height(rbtree));

    rbtree_free(rbtree,NULL);

    return 0;
}