Red-Black Tree

rbtree.c

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

/*
 *
 * 1: every node is either red or black
 * 2: root is black
 * 3: every leaf (NIL) is black
 * 4: if a node is red then both its children are black
 * 5: For each node, every simple path from the node to
 *    descendent leaves contain the same number of black
 *    nodes
 *
 * Black-height of a node: the number of black nodes
 * on any sample path from the node to the leaf
 * */
typedef struct Node {
    struct Node* left;
    struct Node* right;
    struct Node* parent;
    int value;
    int color;
} Node;

Node g_Nil = { NULL, NULL, &g_Nil, 0, 0 };

Node* MakeNode(int value)
{
    Node* node = (Node*)malloc(sizeof(Node));
    node->color = 0;
    node->value = value;
    node->left = &g_Nil;
    node->right = &g_Nil;
    node->parent = &g_Nil;
    return node;
}

// pg 313 Intro to Algorithms 3rd ed.
void LeftRotate(Node** root, Node* x)
{
    // x's left does not change
    // y's right does not change
    assert(root && *root);
    assert(x);

    Node* y = x->right;
    assert(y != &g_Nil);

    // make x the left xhild of y
    if (y->left != &g_Nil)
        y->left->parent = x;

    y->parent = x->parent;

    if (x->parent == &g_Nil)
        *root = y;
    else if (x->parent->left == x)
        x->parent->left = y;
    else
        x->parent->right = y;

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

void RightRotate(Node** root, Node* y)
{
    // x's left does not change
    // y's right does not change
    assert(root && *root);
    assert(y);

    Node* x = y->left;
    assert(x != &g_Nil);

    if (x->right != &g_Nil)
        x->right->parent = y;

    x->parent = y->parent;

    if (y->parent == &g_Nil)
        *root = x;
    else if (y->parent->left == y)
        y->parent->left = x;
    else
        y->parent->right = x;

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

/*
 * After Insert finishes, the red-black properties must hold:
 *
 * 1: every node is either red or black
 * 2: root is black
 * 3: every leaf (NIL) is black
 * 4: if a node is red then both its children are black
 * 5: For each node, every simple path from the node to
 *    to the leaves has the same number of black nodes.
 *
 * When InsertFix up is called, Properties 2 and 4
 * an be violated because z is
 * set to red by Insert. Property 2 is violated if z
 * was set as the root. Property 4 is violated if z's
 * parent is red. Properties 1, 3, and 5 are not violated
 * when InsertFixup is called.
 *
 * At every iteration of the while loop, this three-part
 * invariant is maintained:
 * 1: Node z is red
 * 2: If z.p is the root, then z.p is black
 * 3: If the tree violates any of the red-black properties,
 *    then only poperties 2 and 4 are violated.
 *
 * There are 6 cases in the while loop, with 3 of the cases
 * being symmetric to the other three depending on if z's
 * parent is a left/right child of z's grandparent.
 *
 * Consider the three cases where the red-black properties
 * are violated when z's parent is a left child of z's grandparent.
 * If z's parent is the root, which is black, then the while loop
 * is not entered. That means when the while loop is entered, z's
 * grandparent exists..
 *
 * case 1: z's uncle is red
 *
 *                     C black
 *          A red                     D red
 *      s1            B red        s4        s5
 *                 s2        s3
 *
 * s1, s2, s3, s4 and s5 are subtrees. c is the root, B is
 * the newly inserted node z. y is z's unclude.
 *
 * Property four is violated because A is red and
 * A's right child, B, is red.
 *
 * z's parent, A, is changed to black.
 * z's uncle, D, is changed to black.
 * z's grandparent, C, is changed to red.
 * z is then changed to point to C.
 *
 * At the end of this iteration of the while loop,
 * 1: Node z is red
 * 2: If z's parent is the root, then z's parent is black.
 *     * z's parent (previously z.p.p.p) was not changed.
 * 3: If the tree violates any of the red-black properties,
 *    then only poperties 2 and 4 are violated.
 *    * every node is still either red or black (1), root
 *      is still black (3), and original z's parent and uncle
 *      are not both black, maintaining (5).
 *
 * case 2: z's uncle is black and z is a right child
 * case 3: z's uncle is black and z is a left child
 *
 * in all three cases z's grandparent is black because z's
 * parent is red, and property 4 is violated only between z and
 * z's parent.
 *
 * */
void InsertFixup(Node** root, Node* z)
{
    assert(root && *root);
    assert(z);
    // z must be red
    assert(z->color == 1);

    // The loop terminates when z's parent is black
    while (z->parent->color == 1)
    {
        // if z's parent is the left child of z's grandparent
        if (z->parent == z->parent->parent->left)
        {
            // Node y is z's uncle
            Node* y = z->parent->parent->right;
            // Case 1: z's uncle is red
            if (y->color == 1)
            {
                z->parent->color = 0;
                y->color = 0;
                z->parent->parent->color = 1;
                z = z->parent->parent;
            }
            else
            {
                // Case 2: z's uncle is black and z is a right child
                if (z == z->parent->right)
                {
                    z = z->parent;
                    LeftRotate(root, z);
                }

                // After the above rotation, (or not):
                // Case 3: z's uncle is black and z is a left child
                z->parent->color = 0;
                z->parent->parent->color = 1;
                RightRotate(root, z->parent->parent);
            }
        }
        else // the below is symmetricto above with left/right swapped
        {
            Node* y = z->parent->parent->left;
            if (y->color ==1)
            {
                z->parent->color = 0;
                y->color = 0;
                z->parent->parent->color = 1;
                z = z->parent->parent;
            }
            else
            {
                if (z == z->parent->left)
                {
                    z = z->parent;
                    RightRotate(root, z);
                }

                z->parent->color = 0;
                z->parent->parent->color = 1;
                LeftRotate(root, z->parent->parent);
            }
        }
    }

    (*root)->color = 0;
}

void Insert(Node** root, Node* z)
{
    assert(root && *root);
    assert(z);
    Node* y = &g_Nil;
    Node* x = *root;

    while (x != &g_Nil) {
        y = x;
        if (z->value < x->value)
            x = x->left;
        else
            x = x->right;
    }

    z->parent = y;

    if (y == &g_Nil)
        *root = z;
    else if (z->value < z->parent->value)
        z->parent->left = z;
    else
        z->parent->right = z;

    z->left = &g_Nil;
    z->right = &g_Nil;
    z->color = 1;
    InsertFixup(root, z);
}

/*
 * Places Node v in u's position in the tree,
 * leaving Node u untouched.
 * */
void Transplant(Node** root, Node* u, Node* v)
{
    assert(root);
    assert(*root);
    assert(u);
    assert(v);

    if (u->parent == &g_Nil)
        *root = v;
    else if (u == u->parent->left)
        u->parent->left = v;
    else
        u->parent->right = v;

    v->parent = u->parent;
}

/*
 * Finds the minimum starting from node z
 * */
Node* Minimum(Node* z)
{
    assert(z);
    while (z->left != & g_Nil)
        z = z->left;

    return z;
}

void DeleteFixup(Node** root, Node* x)
{
    assert(root);
    assert(*root);
    assert(x);

    while (x != *root && x->color == 0) {
        // if x is the left child
        if (x == x->parent->left) {

            // w is x's sibling
            Node* w = x->parent->right;
            if (w->color == 1) {
                w->color = 0;
                x->parent->color = 1;
                LeftRotate(root, x->parent);
                w = x->parent->right;
            }

            if (w->left->color == 0 && w->right->color == 0) {
                w->color = 1;
                x = x->parent;
            } else {

                if (w->right->color == 0) {
                    w->left->color = 0;
                    w->color = 1;
                    RightRotate(root, w);
                    w = x->parent->right;
                }

                w->color = x->parent->color;
                x->parent->color = 0;
                w->right->color = 0;
                LeftRotate(root, x->parent);
                x = *root;
            }

        } else { // x is the right child

            // w is x's sibling
            Node* w = x->parent->left;
            if (w->color == 1) {
                w->color = 0;
                x->parent->color = 1;
                RightRotate(root, x->parent);
                w = x->parent->left;
            }

            if (w->left->color == 0 && w->right->color == 0) {
                w->color = 1;
                x = x->parent;
            } else {

                if (w->left->color == 0) {
                    w->right->color = 0;
                    w->color = 1;
                    LeftRotate(root, w);
                    w = x->parent->left;
                }

                w->color = x->parent->color;
                x->parent->color = 0;
                w->left->color = 0;
                RightRotate(root, x->parent);
                x = *root;
            }
        }
    }

    x->color = 0;
}

void Delete(Node** root, Node* z)
{
    assert(root);
    assert(*root);
    assert(z);

    // y is the node removed from the tree or moved within the tree
    Node* y = z;
    // save y's original color each time it is assigned
    // for testing at the end of this function.
    int yOriginalColor = y->color;

    // x is used to keep track of moves into node y's
    // original position
    Node* x = NULL;

    // if y has fewer than two children, y is removed
    if (z->left == &g_Nil) {
        x = z->right;
        Transplant(root, z, z->right);
    } else if (z->right == &g_Nil) {
        x = z->left;
        Transplant(root, z, z->left);
    } else {
        // z has two children
        // y will move into z's original position in the tree
        y = Minimum(z->right);
        yOriginalColor = y->color;
        x = y->right;

        if (y->parent == z) {
            // z's right was the minimum on z's right subtree
            // x is y's right, and that is the same as z's right
            x->parent = y;
        } else {
            // move y's right into the position of y
            Transplant(root, y, y->right);
            y->right = z->right;
            y->right->parent = y;
        }

        // Put y in z's position in the tree
        // then update y and its children with z's data
        Transplant(root, z, y);
        y->left = z->left;
        y->left->parent = y;
        y->color = z->color;
    }

    if (yOriginalColor == 0)
        DeleteFixup(root, x);
}

void FreeTree(Node* root)
{
    if (root == NULL || root == &g_Nil)
        return;

    Node* left = root->left;
    Node* right = root->right;

    printf("Freeing %d\n", root->value);
    free(root);
    FreeTree(left);
    FreeTree(right);
}

void PrintNodesAsDot(Node* z)
{
    assert(z);

    printf("%d [color=%s];\n", z->value, z->color == 0 ? "black" : "red");

    if (z->left != &g_Nil) {
        printf("%d -> %d;\n", z->value, z->left->value);
        PrintNodesAsDot(z->left);
    }

    if (z->right != &g_Nil) {
        printf("%d -> %d;\n", z->value, z->right->value);
        PrintNodesAsDot(z->right);
    }
}

void PrintTreeAsDot(Node* z)
{
    assert(z);
    printf("digraph G {\n");
    PrintNodesAsDot(z);
    printf("}\n");
}

int main(int argc, char** argv)
{
    // Make a tree of all black nodes to test rotation
    Node* root = MakeNode(7);
    root->left = MakeNode(4);
    root->left->parent = root;
    root->left->left = MakeNode(3);
    root->left->left->parent = root->left;
    root->left->left->left = MakeNode(2);
    root->left->left->left->parent = root->left->left;
    root->left->right = MakeNode(6);
    root->left->right->parent = root->left;
    root->right = MakeNode(11);
    root->right->parent = root;
    root->right->left = MakeNode(9);
    root->right->left->parent = root->right;
    root->right->right = MakeNode(18);
    root->right->right->parent = root->right;
    root->right->right->left = MakeNode(14);
    root->right->right->left->parent = root->right->right;
    root->right->right->left->left = MakeNode(12);
    root->right->right->left->left->parent = root->right->right->left;
    root->right->right->left->right = MakeNode(17);
    root->right->right->left->right->parent = root->right->right->left;
    root->right->right->right = MakeNode(19);
    root->right->right->right->parent = root->right->right;
    root->right->right->right->right = MakeNode(22);
    root->right->right->right->right->parent = root->right->right->right;
    root->right->right->right->right->left = MakeNode(20);
    root->right->right->right->right->left->parent = root->right->right->right->right;

    assert(root->right->value == 11);
    assert(root->right->left->value == 9);
    assert(root->right->right->value == 18);
    assert(root->right->right->left->value == 14);
    assert(root->right->right->right->value == 19);

    LeftRotate(&root, root->right);
    assert(root->right->value == 18);
    assert(root->right->left->value == 11);
    assert(root->right->right->value == 19);
    assert(root->right->left->left->value == 9);
    assert(root->right->left->right->value == 14);

    RightRotate(&root, root->right);
    assert(root->right->value == 11);
    assert(root->right->left->value == 9);
    assert(root->right->right->value == 18);
    assert(root->right->right->left->value == 14);
    assert(root->right->right->right->value == 19);

    FreeTree(root);

    // Test insertion
    root = NULL;
    int values[] = {2, 3, 4, 6, 7, 9, 11, 12, 14, 17, 18, 19, 20, 22};
    for (int i = 0; i < sizeof(values)/sizeof(int); ++i)
    {
        if (root == NULL)
            root = MakeNode(values[i]);
        else
            Insert(&root, MakeNode(values[i]));
    }

    PrintTreeAsDot(root);

    // Test deletion

    Node* delete = root->right->right;
    printf("Deleting %d\n", delete->value);
    Delete(&root, delete);
    free(delete);

    delete = root;
    printf("Deleting %d\n", delete->value);
    Delete(&root, delete);
    free(delete);

    delete = root->right->left;
    printf("Deleting %d\n", delete->value);
    Delete(&root, delete);
    free(delete);

    delete = root->left->right;
    printf("Deleting %d\n", delete->value);
    Delete(&root, delete);
    free(delete);

    delete = root->left->left;
    printf("Deleting %d\n", delete->value);
    Delete(&root, delete);
    free(delete);

    printf("\n\n");
    PrintTreeAsDot(root);

    FreeTree(root);

    return 0;
}