kehuadong

#include "common.h"

typedef struct rb_node_t rb_node_t;

struct rb_node_t
{
    rb_node_t*  m_parent;
    rb_node_t*  m_left;
    rb_node_t*  m_right;
    bool        m_red;
    int         m_value;
};

rb_node_t* create_node(rb_node_t* parent, int value)
{
    rb_node_t* p    = (rb_node_t*)calloc(1, sizeof(rb_node_t));
    p->m_parent     = parent;
    p->m_red        = true;
    p->m_value      = value;
    return p;
}

static void rotate_left(rb_node_t** root, rb_node_t* p)
{
    rb_node_t* pp   = p->m_parent;
    rb_node_t* r    = p->m_right;

    p->m_right = r->m_left;
    r->m_left = p;
    if (pp == NULL)
    {
        *root = r;
    }
    else if (pp->m_left == p)
    {
        pp->m_left = r;
    }
    else
    {
        pp->m_right = r;
    }

    if (p->m_right)
    {
        p->m_right->m_parent = p;
    }
    p->m_parent = r;
    r->m_parent = pp;
}

static void rotate_right(rb_node_t** root, rb_node_t* p)
{
    rb_node_t* pp   = p->m_parent;
    rb_node_t* l    = p->m_left;

    p->m_left = l->m_right;
    l->m_right = p;
    if (pp == NULL)
    {
        *root = l;
    }
    else if (pp->m_left == p)
    {
        pp->m_left = l;
    }
    else
    {
        pp->m_right = l;
    }
    
    if (p->m_left)
    {
        p->m_left->m_parent = p;
    }
    p->m_parent = l;
    l->m_parent = pp;
}

static void fix_after_insert(rb_node_t** root, rb_node_t* p)
{
    while (p->m_red)
    {
        rb_node_t* pp = p->m_parent;
        rb_node_t* b = (pp->m_left == p) ? pp->m_right : pp->m_left;
        if (b && b->m_red)
        {
            p->m_red = false;
            b->m_red = false;

            if (pp->m_parent)
            {
                pp->m_red = true;
                p = p->m_parent;
                continue;
            }
        }
        else
        {
            if (p == pp->m_left)
            {
                if (p->m_right && p->m_right->m_red)
                {
                    rotate_left(root, p);
                    p = p->m_parent;
                }

                p->m_red = false;
                p->m_parent->m_red = true;
                rotate_right(root, p->m_parent);
            }
            else
            {
                if (p->m_left && p->m_left->m_red)
                {
                    rotate_right(root, p);
                    p = p->m_parent;
                }

                p->m_red = false;
                p->m_parent->m_red = true;
                rotate_left(root, p->m_parent);
            }
        }
        break;
    }
}

void rb_insert(rb_node_t** root, int value)
{
    if (*root == NULL)
    {
        *root = create_node(NULL, value);
        (*root)->m_red = false;
        return;
    }

    rb_node_t* c = *root;
    rb_node_t* p = NULL;
    while (c)
    {
        p = c;
        if (value <= c->m_value)
        {
            c = c->m_left;
        }
        else
        {
            c = c->m_right;
        }
    }

    if (value <= p->m_value)
    {
        p->m_left = create_node(p, value);
    }
    else
    {
        p->m_right = create_node(p, value);
    }

    fix_after_insert(root, p);
}

int rb_depth(rb_node_t* p)
{
    if (p == NULL)
    {
        return 0;
    }

    int a = rb_depth(p->m_left);
    int b = rb_depth(p->m_right);

    return 1 + (a < b ? b : a);
}

int rb_count(rb_node_t* p)
{
    if (p == NULL)
    {
        return 0;
    }
    return 1 + rb_count(p->m_left) + rb_count(p->m_right);
}


int main()
{
    rb_node_t* root = NULL;
    for (int i = 0; i < 1024*1024; i++)
    {
        rb_insert(&root, i);
    }
    printf("rb_count = %d, rb_depth = %d\n", rb_count(root), rb_depth(root));
}