AVL trees

avl.c

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

#define AVL_AS_LIST

typedef int AvlKey;

struct AvlNode {
    AvlKey key;
    char balance;
    struct AvlNode *left;
    struct AvlNode *right;
#ifdef AVL_AS_LIST
    struct AvlNode *prev;
    struct AvlNode *next;
#endif
};

typedef struct AvlNode AvlNode;

AvlNode* avl_insert(AvlNode* tree, AvlKey key);
AvlNode* avl_delete(AvlNode* tree, AvlKey key);
AvlNode* avl_upper(AvlNode* tree, AvlKey key);
AvlNode* avl_lower(AvlNode* tree, AvlKey key);

int avl_cmp(AvlKey p, AvlKey q) {
    return p - q;
}

#ifdef AVL_AS_LIST
static
void avl_list_insert_before(AvlNode* list, AvlNode* node) {
    if (list->prev) list->prev->next = node;
    node->prev = list->prev;
    node->next = list;
    list->prev = node;
}

static
void avl_list_insert_after(AvlNode* list, AvlNode* node) {
    if (list->next) list->next->prev = node;
    node->prev = list;
    node->next = list->next;
    list->next = node;
}

static
void avl_list_delete(AvlNode* node) {
    if (node->prev) node->prev->next = node->next;
    if (node->next) node->next->prev = node->prev;
    node->prev = NULL;
    node->next = NULL;
}
#endif

static
int avl_height(AvlNode* tree) {
    if (tree == NULL) {
        return 0;
    } else {
        int h1, h2;
        h1 = avl_height(tree->left);
        h2 = avl_height(tree->right);
        if (h1 > h2) {
            return 1 + h1;
        } else {
            return 1 + h2;
        }
    }
}

static
AvlNode* avl_lowest(AvlNode* node) {
    if (node == NULL) return NULL;
    if (node->left) return avl_lowest(node->left);
    return node;
}

static
AvlNode* avl_highest(AvlNode* node) {
    if (node == NULL) return NULL;
    if (node->right) return avl_highest(node->right);
    return node;
}

#define MAX(a, b) ((a) > (b) ? (a) : (b))

static
AvlNode* avl_rotate(AvlNode* parent, AvlNode* child) {
    if (child == parent->left) {
        parent->left = child->right;
        child->right = parent;
        parent->balance += MAX(0, -1 * child->balance) + 1;
        child->balance += MAX(0, parent->balance) + 1;
    } else if (child == parent->right) {
        parent->right = child->left;
        child->left = parent;
        parent->balance -= MAX(0, child->balance) + 1;
        child->balance -= MAX(0, -1 * parent->balance) + 1;
    }
    return child;
}

static
AvlNode* avl_rebalance_left(AvlNode* tree) {
    if (tree->left->balance == -1) {
        return avl_rotate(tree, tree->left);
    } else if (tree->left->balance == 1) {
        tree->left = avl_rotate(tree->left, tree->left->right);
        return avl_rotate(tree, tree->left);
    } else {
        return avl_rotate(tree, tree->left);
    }
}

static
AvlNode* avl_rebalance_right(AvlNode* tree) {
    if (tree->right->balance == 1) {
        return avl_rotate(tree, tree->right);
    } else if (tree->right->balance == -1) {
        tree->right = avl_rotate(tree->right, tree->right->left);
        return avl_rotate(tree, tree->right);
    } else {
        return avl_rotate(tree, tree->right);
    }
}

AvlNode* avl_insert(AvlNode* tree, AvlKey key) {
    if (tree == NULL) {
        tree = malloc(sizeof(AvlNode));
        tree->key = key;
        tree->balance = 0;
        tree->left = NULL;
        tree->right = NULL;
#ifdef AVL_AS_LIST
        tree->prev = NULL;
        tree->next = NULL;
#endif
    } else {
        if (avl_cmp(key, tree->key) < 0) {
            if (tree->left == NULL) {
                tree->balance = (tree->right ? 0 : -1);
                tree->left = avl_insert(tree->left, key);
#ifdef AVL_AS_LIST
                avl_list_insert_before(tree, tree->left);
#endif
            } else {
                char left_was_zero = tree->left->balance == 0;
                tree->left = avl_insert(tree->left, key);
                if (left_was_zero && tree->left->balance != 0) {
                    tree->balance--;
                    if (tree->balance == -2) {
                        tree = avl_rebalance_left(tree);
                    }
                }
            }
        } else if (avl_cmp(key, tree->key) > 0) {
            if (tree->right == NULL) {
                tree->balance = (tree->left ? 0 : 1);
                tree->right = avl_insert(tree->right, key);
#ifdef AVL_AS_LIST
                avl_list_insert_after(tree, tree->right);
#endif
            } else {
                char right_was_zero = tree->right->balance == 0;
                tree->right = avl_insert(tree->right, key);
                if (right_was_zero && tree->right->balance != 0) {
                    tree->balance++;
                    if (tree->balance == 2) {
                        tree = avl_rebalance_right(tree);
                    }
                }
            }
        }
    }
    return tree;
}

AvlNode* avl_delete(AvlNode* tree, AvlKey key) {
    if (tree == NULL) return tree;
    if (avl_cmp(key, tree->key) < 0) {
        char had_left = tree->left != NULL;
        char left_was_nonzero = tree->left && tree->left->balance != 0;
        tree->left = avl_delete(tree->left, key);
        if ((had_left && tree->left == NULL) || (left_was_nonzero && tree->left->balance == 0)) {
            tree->balance++;
        }
        if (tree->balance == 2) {
            tree = avl_rebalance_right(tree);
        }
        return tree;
    } else if (avl_cmp(key, tree->key) > 0) {
        char had_right = tree->right != NULL;
        char right_was_nonzero = tree->right && tree->right->balance != 0;
        tree->right = avl_delete(tree->right, key);
        if ((had_right && tree->right == NULL) || (right_was_nonzero && tree->right->balance == 0)) {
            tree->balance--;
        }
        if (tree->balance == -2) {
            tree = avl_rebalance_left(tree);
        }
        return tree;
    }
    // key == tree->key
    if (tree->left && tree->right) {
#ifdef AVL_AS_LIST
        AvlKey tmp = tree->prev->key;
#else
        AvlKey tmp = avl_highest(tree->left)->key;
#endif
        AvlNode* new_tree = avl_delete(tree, tmp);
        tree->key = tmp;
        return new_tree;
    }
    AvlNode* ret = NULL;
    if (tree->left) {
        ret = tree->left;
    } else if (tree->right) {
        ret = tree->right;
    }
#ifdef AVL_AS_LIST
    avl_list_delete(tree);
#endif
    free(tree);
    return ret;
}

static
AvlNode* avl_find_upper(AvlNode* tree, AvlKey key, AvlNode* candidate) {
    if (tree == NULL) return candidate;
    if (avl_cmp(key, tree->key) == 0) {
        return tree;
    } else if (avl_cmp(key, tree->key) < 0) {
        return avl_find_upper(tree->left, key, tree);
    } else {
        return avl_find_upper(tree->right, key, candidate);
    }
}

AvlNode* avl_upper(AvlNode* tree, AvlKey key) {
    return avl_find_upper(tree, key, NULL);
}

static
AvlNode* avl_find_lower(AvlNode* tree, AvlKey key, AvlNode* candidate) {
    if (tree == NULL) return candidate;
    if (avl_cmp(key, tree->key) == 0) {
        return tree;
    } else if (avl_cmp(key, tree->key) < 0) {
        return avl_find_lower(tree->left, key, candidate);
    } else {
        return avl_find_lower(tree->right, key, tree);
    }
}

AvlNode* avl_lower(AvlNode* tree, AvlKey key) {
    return avl_find_lower(tree, key, NULL);
}

int
avl_check_balance(AvlNode* tree) {
    if (tree == NULL) return 0;
    int left_height = avl_check_balance(tree->left);
    int right_height = avl_check_balance(tree->right);
    int correct_balance = right_height - left_height;
    if (correct_balance != tree->balance || tree->balance < -1 || tree->balance > 1) {
        printf("incorrect balance %d, should be %d\n", tree->balance, correct_balance);
        exit(1);
    }
    return 1 + MAX(left_height, right_height);
}

test.c

#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include "avl.c"

#define RANGE 50000

int main() {
    int i;
    AvlNode* tree = NULL;

    srand(time(NULL));

    for (i = 0; i < RANGE; i++) {
        tree = avl_insert(tree, rand() % RANGE);
        //avl_check_balance(tree);
    }
    for (i = 0; i < RANGE; i++) {
        tree = avl_delete(tree, rand() % RANGE);
        //avl_check_balance(tree);
    }

    return 0;
}