[C] RB Tree 구현

 

rb_tree.h

/* rb_tree.h */

#ifndef _RB_TREE_H_
#define _RB_TREE_H_

#include <stdbool.h>

typedef struct rb_tree_t rb_tree_t;
typedef struct rb_node_t rb_node_t;
typedef enum rb_color_e rb_color_e;

enum rb_color_e {
    RED,
    BLACK,
};

struct rb_node_t {
    void* data;

    struct rb_node_t* left;
    struct rb_node_t* right;
    struct rb_node_t* parent;
    enum rb_color_e color;
};

struct rb_tree_t {
    struct rb_node_t* root;
    struct rb_node_t* leaf;

    /*
     * if (a > b)
     *  return true;
     * else
     *  return false;
     */
    bool (*larger_than)(void* a, void* b);

    /*
     * if (a==b)
     *  return true;
     * else 
     * return false;
     */
    bool (*equal)(void* a, void* b);
};

void rb_tree_init(rb_tree_t* tree, bool (*larger_than)(void* a, void* b), bool (*equal)(void* a, void* b));
rb_node_t* rb_exist(rb_tree_t* tree, void* data);
void rb_insert(rb_tree_t* tree, void * new_data);
bool rb_delete(rb_tree_t* tree, void* data);
void rb_destroy(rb_tree_t* tree);

void rb_for_inorder(rb_tree_t* tree, void (*probe)(void *data));
void rb_for_preorder(rb_tree_t* tree, void (*probe)(void* data));
void rb_for_postorder(rb_tree_t* tree, void (*probe)(void* data));

#endif

 

rb_tree.c

/* rb_tree.c */

#include <stdlib.h>

#include "rb_tree.h"

rb_node_t* __rb_min(rb_tree_t *tree, rb_node_t* node)
{
    while (node->left != tree->leaf)
        node = node->left;

    return node;
}

rb_node_t *rb_exist(rb_tree_t *tree, void* data)
{
    rb_node_t* x = tree->root;
    rb_node_t* p = NULL;

    if (tree->root == tree->leaf)
        return NULL;

    while (x != NULL && !tree->equal(data, x->data)) {
        p = x;

        if (tree->larger_than(p->data, data))
            x = p->left;
        else
            x = p->right;

        if (x == tree->leaf)
            return NULL;
    }

    return x;
}

static rb_node_t* alloc_rb_node(void)
{
    rb_node_t* ret;
    ret = (rb_node_t*)malloc(sizeof(rb_node_t));

    if (!ret)
        return NULL;

    ret->color = BLACK;
    ret->parent = NULL;
    ret->left = NULL;
    ret->right = NULL;

    return ret;
}

static void free_rb_node(rb_node_t *node)
{
    free(node);
}

static void rb_rotate_left(rb_tree_t *tree, rb_node_t* x)
{
    rb_node_t* y;

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

    if (y->left != tree->leaf)
        y->left->parent = x;

    y->parent = x->parent;

    if (!x->parent)
        tree->root = y;
    else if (x == x->parent->left)
        x->parent->left = y;
    else
        x->parent->right = y;

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

static void rb_rotate_right(rb_tree_t *tree, rb_node_t* y)
{
    rb_node_t* x;
    
    x = y->left;
    y->left = x->right;

    if (x->right != tree->leaf)
        x->right->parent = y;

    if (!y->parent)
        tree->root = x;
    else if (y == y->parent->left)
        y->parent->left = x;
    else
        y->parent->right = x;

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

static void rb_insert_fixup(rb_tree_t *tree, rb_node_t* n)
{
    while (n != tree->root && n->parent->color == RED) {
        rb_node_t* p = n->parent;
        rb_node_t* g = p->parent;
        rb_node_t* u = (p == g->left) ? g->right : g->left;
        bool is_p_left = (p == g->left) ? true : false;

        if (u && u->color == RED) {
            n->parent->color = BLACK;
            u->color = BLACK;
            g->color = RED;
            n = g;
        } else {
            if (n == (is_p_left ? p->right : p->left)) {
                /*
                 * p==left && n==right
                 * p==right && n==left
                 */
                n = n->parent;
                if (is_p_left)
                    rb_rotate_left(tree, n);
                else
                    rb_rotate_right(tree, n);
            }

            n->parent->color = BLACK;
            g->color = RED;
            if (is_p_left)
                rb_rotate_right(tree, g);
            else
                rb_rotate_left(tree, g);
        }
    }
    tree->root->color = BLACK;
}

void rb_insert(rb_tree_t *tree, void* new_data)
{
    rb_node_t* x = tree->root;
    rb_node_t* p = NULL;
    rb_node_t* n = alloc_rb_node();

    n->color = RED;
    n->left = tree->leaf;
    n->right = tree->leaf;
    n->data = new_data;

    while (x != tree->leaf) {
        p = x;
        if (tree->larger_than(x->data, n->data))
            x = x->left;
        else
            x = x->right;
    }

    n->parent = p;
    
    if (p == NULL)
        tree->root = n;
    else if (tree->larger_than(n->data, p->data))
        p->right = n;
    else
        p->left = n;

    if (n == tree->root) {
        n->color = BLACK;
        return;
    }

    if (n->parent == tree->root)
        return;

    rb_insert_fixup(tree, n);
}

/* put 'v' to location of 'u' */
static void rb_trans_plant(rb_tree_t *tree, rb_node_t* u, rb_node_t* v)
{
    if (u == tree->root)
        tree->root = v;
    else if (u == u->parent->left)
        u->parent->left = v;
    else
        u->parent->right = v;

    v->parent = u->parent;
}


static void rb_delete_fixup(rb_tree_t *tree, rb_node_t* x)
{
    rb_node_t* b;

    while (x != tree->root && x->color == BLACK) {
        if (x == x->parent->left) {
            b = x->parent->right;

            if (b->color == RED) {
                b->color = BLACK;
                x->parent->color = RED;
                rb_rotate_left(tree, x->parent);
                b = x->parent->right;
            }

            if (b->left->color == BLACK && b->right->color == BLACK) {
                b->color = RED;
                x = x->parent;
            } else {
                if (b->right->color == BLACK) {
                    b->left->color = BLACK;
                    b->color = RED;
                    rb_rotate_right(tree, b);
                    b = x->parent->right;
                }

                b->color = x->parent->color;
                x->parent->color = BLACK;
                b->right->color = BLACK;
                rb_rotate_left(tree, x->parent);
                x = tree->root;
            }

        } else {
            b = x->parent->left;

            if (b->color == RED) {
                b->color = BLACK;
                x->parent->color = RED;
                rb_rotate_right(tree, x->parent);
                b = x->parent->left;
            }

            if (b->left->color == BLACK && b->right->color == BLACK) {
                b->color = RED;
                b = b->parent;
            } else {
                if (b->left->color == BLACK) {
                    b->right->color = BLACK;
                    b->color = RED;
                    rb_rotate_left(tree, b);
                    b = x->parent->left;
                }

                b->color = x->parent->color;
                x->parent->color = BLACK;
                b->left->color = BLACK;
                rb_rotate_right(tree, x->parent);
                x = tree->root;
            }
        }
    }
    x->color = BLACK;
    tree->root->color = BLACK;
}

bool rb_delete(rb_tree_t *tree, void *data)
{
    rb_node_t* t = rb_exist(tree, data);
    rb_node_t* x;
    rb_node_t* y;
    rb_color_e origin_corlor;

    if (!t)
        return false;

    origin_corlor = t->color;

    if (t->left == tree->leaf) {
        x = t->right;
        rb_trans_plant(tree, t, t->right);
    } else if (t->right == tree->leaf) {
        x = t->left;
        rb_trans_plant(tree, t, t->left);
    } else {
        y = __rb_min(tree, t->right);
        origin_corlor = y->color;

        /* child of y is leaf */
        x = y->right;

        if (y->parent == t)
            x->parent = y;
        else {
            rb_trans_plant(tree, y, y->right);
            y->right = t->right;
            y->right->parent = y;
        }
        rb_trans_plant(tree, t, y);
        y->left = t->left;
        y->left->parent = y;
        y->color = t->color;
    }

    free_rb_node(t);

    if (origin_corlor == BLACK) {
        rb_delete_fixup(tree, x);
    }

    return true;
}

static void __rb_destroy(rb_tree_t* tree, rb_node_t* node)
{
    if (node == tree->leaf)
        return;

    __rb_destroy(tree, node->left);
    __rb_destroy(tree, node->right);
    free(node);
}

void rb_destroy(rb_tree_t *tree)
{
    rb_node_t* node = tree->root;

    __rb_destroy(tree, node);
}

static void __rb_for_inorder(rb_tree_t* tree, rb_node_t* node, void (*probe)(void* data))
{
    if (node == tree->leaf)
        return;

    __rb_for_inorder(tree, node->left, probe);
    probe(node->data);
    __rb_for_inorder(tree, node->right, probe);
}

void rb_for_inorder(rb_tree_t *tree, void (*probe)(void* data))
{
    rb_node_t* node = tree->root;
    __rb_for_inorder(tree, node, probe);
}

static void __rb_for_preorder(rb_tree_t* tree, rb_node_t* node, void (*probe)(void* data))
{
    if (node == tree->leaf)
        return;

    probe(node->data);
    __rb_for_preorder(tree, node->left, probe);
    __rb_for_preorder(tree, node->right, probe);
}

void rb_for_preorder(rb_tree_t* tree, void (*probe)(void* data))
{
    rb_node_t* node = tree->root;
    __rb_for_preorder(tree, node, probe);
}

static void __rb_for_postorder(rb_tree_t* tree, rb_node_t* node, void (*probe)(void* data))
{
    if (node == tree->leaf)
        return;

    __rb_for_postorder(tree, node->left, probe);
    __rb_for_postorder(tree, node->right, probe);
    probe(node->data);
}

void rb_for_postorder(rb_tree_t* tree, void (*probe)(void* data))
{
    rb_node_t* node = tree->root;
    __rb_for_postorder(tree, node, probe);
}

void rb_tree_init(rb_tree_t* tree, bool (*larger_than)(void* a, void* b), bool (*equal)(void* a, void* b))
{
    tree->leaf = alloc_rb_node();
    tree->root = tree->leaf;

    tree->larger_than = larger_than;
    tree->equal = equal;
}

 

main.c

/* main.c */
#include <stdio.h>
#include <stdlib.h>

#include "rb_tree.h"

/* user defined data and operator */
struct data_t {
    int key;
};

struct data_t* alloc_data(void)
{
    return (struct data_t*)malloc(sizeof(struct data_t));
}

void free_data(void* data)
{
    free(data);
}

bool data_larger_than(void* a, void* b)
{
    if (((struct data_t *)a)->key > ((struct data_t*)b)->key)
        return true;

    return false;
}

bool data_equal(void* a, void* b)
{
    if (((struct data_t*)a)->key == ((struct data_t*)b)->key)
        return true;

    return false;
}

void data_print(void* d)
{
    struct data_t* data = (struct data_t *)d;
    printf("%d ", data->key);
}

/* main */
enum menu_e {
    INSERT = 1,
    DELETE = 2,
    PRINT_IN_ORDER = 3,
    EXIT = 99,
};

void print_main_menu(void)
{
    printf("::: MENU :::\n");
    printf("1. insert\n");
    printf("2. delete\n");
    printf("3. print(in order)\n");
    printf("------------\n");
    printf("99. exit\n");
    printf("------------\n");
}

int main()
{
    struct rb_tree_t rb_tree;
    rb_tree_init(&rb_tree, data_larger_than, data_equal);

    while (true) {
        int select;

        print_main_menu();
        scanf("%d", &select);

        switch (select)
        {
        case INSERT:
        {
            int key;
            printf("input value : ");
            scanf("%d", &key);
            struct data_t* data = alloc_data();
            if (data) {
                data->key = key;

                if (!rb_exist(&rb_tree, data))
                    rb_insert(&rb_tree, data);
                else {
                    printf("ERROR: %d alreday exist\n", key);
                    free_data(data);
                }
            } else
                printf("memory alloc fail\n");
        }
        break;
        case DELETE:
        {
            int key;
            printf("input value : ");
            scanf("%d", &key);
            struct data_t* data = alloc_data();
            if (data) {
                data->key = key;
                if (!rb_delete(&rb_tree, data))
                    printf("ERROR: %d does not exist\n", key);
                free_data(data);
            } else
                printf("memory alloc fail\n");
        }
        break;
        case PRINT_IN_ORDER:
        {
            printf("print in order : \n");
            rb_for_inorder(&rb_tree, data_print);
            printf("\n");
        }
        break;
        case EXIT:
        {
            goto out;
        }
        break;
        default:
            break;
        }
        printf("\n");
    }

out:
    rb_for_postorder(&rb_tree, free_data);
    rb_destroy(&rb_tree);

    return 0;
}

 

---

※ 참고

https://www.programiz.com/dsa/red-black-tree

Red-Black Tree

https://www.cs.usfca.edu/~galles/visualization/RedBlack.html

Red/Black Tree Visualization

.