Saturday, May 10, 2008

do your trees lean left

I haven't found any concrete examples of the algorithms described by Sedwick.

So, below is my crack at it... maybe someday I'll get around to annotating it so this will be meaningful, for now you are on your own.

static int
__tree_node_is_red(tree_node_t *n)
{
    return n && n->color == RED;
}

static int
__tree_node_is_black(tree_node_t *n)
{
    return n && n->color == BLACK;
}

static tree_node_t *
__tree_node_rotate(tree_node_t *n, tree_node_link_t dir)
{
    tree_node_t *p;
    assert(dir == LEFT || dir == RIGHT);

    p = n->link[!dir];

    /* child of p -> child of n */
    n->link[!dir] = p->link[dir];
    if (n->link[!dir])
        n->link[!dir]->link[PARENT] = n;

    /* parent of n -> parent of p */
    p->link[PARENT] = n->link[PARENT];
    if (p->link[PARENT]) {
        if (n == p->link[PARENT]->link[dir]) {
            p->link[PARENT]->link[dir] = p;
        } else {
            p->link[PARENT]->link[!dir] = p;
        }
    }

    /* swap p and n */
    p->link[dir] = n;
    n->link[PARENT] = p;

    return p;
}

static tree_node_t *
__tree_node_split(tree_node_t *n)
{
    tree_node_t *p = __tree_node_rotate(n, RIGHT);
    p->link[LEFT]->color = BLACK;
    return p;
}

static tree_node_t *
__tree_node_lean(tree_node_t *n, tree_node_link_t dir)
{
    tree_node_t *p = __tree_node_rotate(n, dir);
    p->color = p->link[dir]->color;
    p->link[dir]->color = RED;
    return p;
}

void *
tree_insert(tree_t *t, void *key, void *val)
{
    int c;
    tree_node_t *n, *p;
    void *r = NULL;

    assert(t);

    for (p = NULL, n = t->root; n; p = n, n = n->link[c > 0]) {
#ifdef BALANCE
        /* split four nodes on the way down the path */
        if (__tree_node_is_red(n->link[LEFT]))
            if (__tree_node_is_red(n->link[LEFT]->link[LEFT]))
                n = __tree_node_split(n);
#endif

        /* replace an existing entry */
        if ((c = t->cmp(key, n->entry[KEY])) == 0) {
            r = n->entry[VALUE];
            n->entry[VALUE] = val;
            break;
        }
    }

    /* adding a new node */
    try {
        if (!r) {
            n = __tree_node_create(t, p, key, val);
            throw_unless(n) out_of_memory;

            if (p) {
                p->link[c > 0] = n;
            }
        }

#ifdef BALANCE
        /* lean red links left on the way up the path */
        for ( ; p; n = p, p = p->link[PARENT])
            if (__tree_node_is_red(p->link[RIGHT]))
                p = __tree_node_lean(p, LEFT);
#endif

        /* created (or rotated) a new root if parent is null */
        if (!n->link[PARENT]) {
            t->root = n;
            n->color = BLACK;
        }
        t->size++;
    } catch (out_of_memory) { }
    return NULL;
}

static tree_node_t *
__tree_node_move(tree_node_t *n, tree_node_link_t dir)
{
    n->color = BLACK;
    n->link[dir]->color = RED;
    if (__tree_node_is_red(n->link[!dir]->link[LEFT])) {
        if (dir) {
            n = __tree_node_rotate(n, dir);
            n->color = RED;
            n->link[!dir]->color = BLACK;
        } else {
            n->link[!dir] = __tree_node_rotate(n->link[!dir], !dir);
            n = __tree_node_rotate(n, dir);
        }
    } else {
        n->link[!dir]->color = RED;
    }
    return n;
}

void *
tree_remove(tree_t *t, void *key)
{
    int c = 1;
    tree_node_t *n, *p;
    void *r = NULL;

    assert(t);

    for (p = NULL, n = t->root; n && c; p = n, n = n->link[c > 0]) {
        c = t->cmp(key, n->entry[KEY]);
        if (c < 0) {
            /* push red right */
            if (__tree_node_is_black(n->link[LEFT]) &&
                    __tree_node_is_black(n->link[LEFT]->link[LEFT]))
                n = __tree_node_move(n, LEFT);
        } else if (c > 0) {
            if (__tree_node_is_red(n->link[LEFT]))
                n = __tree_node_lean(n, RIGHT);

            if (__tree_node_is_black(n->link[RIGHT]) &&
                    __tree_node_is_black(n->link[RIGHT]->link[LEFT]))
                n = __tree_node_move(n, RIGHT);
        } else {
            if (n->link[RIGHT]) {
                tree_node_t *m = n->link[RIGHT];
                /* find successor */
                while (m->link[LEFT])
                    m = m->link[LEFT];

                /* replace and continue search */
                n->entry[KEY] = m->entry[KEY];
                n->entry[VALUE] = m->entry[VALUE];
                key = m->entry[KEY];
                c = 1; /* force move right */
            } else if (n->link[PARENT]) {
                /* delete leaf, replace parent's pointer to this
                 * node with pointer to this node's left child */
                p->link[p->link[LEFT] != n] = n->link[LEFT];
                if (n->link[LEFT])
                    n->link[LEFT]->link[PARENT] = p;
            } else {
                /* delete root */
                t->root = n->link[LEFT];
                if (n->link[LEFT])
                    n->link[LEFT]->link[PARENT] = NULL;
            }
        }
    }

    for (n = p; n->link[PARENT]; n = n->link[PARENT])
        if (__tree_node_is_red(n->link[RIGHT]))
            n = __tree_node_lean(n, LEFT);

    __tree_node_delete(t, p);

    return r;
}