AVL 樹實現

一個 AVL 樹的 C 實現。先定義接口 avl.h：

#ifndef _AVLTREE_H_
#define _AVLTREE_H_

typedef struct avltree_t avltree_t;

/*
 * Allocates a new empty AVL tree.
 *
 * @return a pointer to the tree
 */
avltree_t* avl_alloc(void);

/*
 * Frees an AVL tree.
 *
 * @param tree  the tree to free.
 */
void avl_free(avltree_t* /*tree*/);

/*
 * Inserts into an AVL tree.
 *
 * @param elem  the element to insert
 */
void avl_insert(avltree_t* /*tree*/, int /*elem*/);

/*
 * Deletes an element from an AVL tree. Does nothing if the element is not in
 * the tree.
 *
 * @param elem  the element to delete
 */
void avl_delete(avltree_t* /*tree*/, int /*elem*/);

/*
 * Searches for an element in an AVL tree.
 *
 * @param tree      the tree to search on
 * @param needle    the element to search
 * @return          1 if `needle' is found or 0 otherwise
 */
int avl_search(avltree_t* /*tree*/, int /*needle*/);

/*
 * Graphs the tree visually using pure text.
 *
 * @param tree  the tree to graph
 */
void avl_graph(avltree_t* /*tree*/);

#endif /* _AVLTREE_H_ */

avl.h 中聲明了一個抽象數據類型 avltree_t，在 avl.c 中首先需要實現的就是這個結構。AVL 樹結構的實現方法變化很多，這裏除了常規的子節點指針和元素鍵之外只用到了一個額外的域保存子樹的高度，實則還可以用一個指針域保存父節點的地址，方便進行再平衡（rebalance）時直接通過域回溯到父節點調整高度並進行旋轉操作。下面的實現由於是通過遞歸提供的隱式棧（函數調用棧）實現回溯，所以並不需要每個節點存儲父節點地址。子樹的高度在實踐時會有一定的上限：理論上在 4GB 的 32 位取址空間中只有 4294967296 的可用字節（這還不計系統、其他程序等所需的資源），若按每個樹節點 16 字節算，最多只可能有 4194304 個節點，所以樹的高度在最好情況下只有 22 左右，而最壞也不會超過 30^[1]。這意味著每個子樹的高度只需要 5 個二進制位就能存儲，在某些空間緊張的場合下可以作為位域與其它域共享一塊整 8 位的內存。

typedef struct node_t node_t;

struct node_t
{
    int elem;
    int height;
    node_t* lchild;
    node_t* rchild;
};

struct avltree_t
{
    node_t* root;
};

基本的內存分配和釋放函數：

avltree_t*
avl_alloc(void)
{
    return calloc(1, sizeof (avltree_t));
}

void
free_nodes(node_t* root)
{
    if (root == 0)
        return;
    free_nodes(root->lchild);
    free_nodes(root->rchild);
    free(root);
}

void
avl_free(avltree_t* tree)
{
    if (tree != 0) {
        free_nodes(tree->root);
        free(tree);
    }
}

一些輔助函數，分別用於計算兩個數的最大值、計算節點高度、更新節點高度以及計算節點平衡因子。

int max(int a, int b)
{
    return a < b ? b : a;
}

/* Empty tree has height -1, tree with one node has height 0, etc. */
int height(node_t* node)
{
    return node != 0 ? node->height : -1;
}

void update_height(node_t* node)
{
    node->height = max(height(node->lchild), height(node->rchild)) + 1;
}

int bfactor(node_t* node)
{
    return height(node->lchild) - height(node->rchild);
}

旋轉操作。使用雙重間接引用是為了避免回溯操縱父節點，因為當沒有父節點引用時只要保存了父節點結構中子節點指針的地址，就可以直接通過間接引用修改父節點所引用的子節點。

/*
 * The following rotation and rebalance operations assume the input trees have
 * a height sufficiently large to perform rotations.
 */

void rotate_left(node_t** node)
{
    node_t* parent = *node;
    node_t* rchild = parent->rchild;
    node_t* rlchild = rchild->lchild;
    rchild->lchild = parent;
    parent->rchild = rlchild;
    *node = rchild;
    update_height(parent);
}

void rotate_right(node_t** node)
{
    node_t* parent = *node;
    node_t* lchild = parent->lchild;
    node_t* lrchild = lchild->rchild;
    lchild->rchild = parent;
    parent->lchild = lrchild;
    *node = lchild;
    update_height(parent);
}

再平衡函数。当平衡因子变为 2 或 -2 时，树不再符合 AVL 树的平衡特性，故而需要通过旋转操作进行再平衡。

void
rebalance(node_t** pnode)
{
    node_t* node = *pnode;
    switch (bfactor(node)) {
    /* Heavy to the left. */
    case 2:
        if (bfactor(node->lchild) == -1)
            rotate_left(&node->lchild);
        rotate_right(pnode);
        break;
    /* Heavy to the right. */
    case -2:
        if (bfactor(node->rchild) == 1)
            rotate_right(&node->rchild);
        rotate_left(pnode);
        break;
    }
}

void
adjust_tree(node_t** pnode)
{
    rebalance(pnode);
    update_height(*pnode);
}

最後是插入、刪除、搜索操作以及一個輸出函數的實現。刪除運算相較插入運算來說略為繁雜，需要考慮的場合更多，可能的再平衡次數也更多，但仍然可在 O(log n) 時間內完成。

void
insert_node(node_t** pnode, int elem)
{
    node_t* node = *pnode;
    if (node == 0) {
        node = *pnode = calloc(1, sizeof (node_t));
        node->elem = elem;
    } else {
        insert_node(elem < node->elem ? &node->lchild : &node->rchild, elem);
        adjust_tree(pnode);
    }
}

void
avl_insert(avltree_t* tree, int elem)
{
     insert_node(&tree->root, elem);
}

/* Removes and returns the leftmost child of a node. */
node_t*
remove_leftmost(node_t** pnode)
{
    node_t* node = *pnode;
    if (node->lchild == 0) {
        *pnode = node->rchild;
    } else {
        node = remove_leftmost(&node->lchild);
        adjust_tree(pnode);
    }
    return node;
}

void
delete_node(node_t** pnode, int elem)
{
    node_t* node = *pnode;
    if (node == 0) {
        return;
    } else if (elem == node->elem) {
        /* `node' is a leaf or has only one child. */
        if (node->lchild == 0) {
            *pnode = node->rchild;
        } else if (node->rchild == 0) {
            *pnode = node->lchild;
        /* `node' has two children. */
        } else {
            node_t* succ = remove_leftmost(&node->rchild);
            succ->lchild = node->lchild;
            succ->rchild = node->rchild;
            /* Replace `node' with its in-order successor. */
            *pnode = succ;
            adjust_tree(pnode);
        }
        free(node);
    } else {
        if (elem < node->elem)
            delete_node(&node->lchild, elem);
        else
            delete_node(&node->rchild, elem);
        adjust_tree(pnode);
    }
}

void
avl_delete(avltree_t* tree, int elem)
{
    delete_node(&tree->root, elem);
}

/*
 * Searches for an AVL tree node and returns a pointer to it if found or 0
 * otherwise.
 */
node_t*
search_node(node_t* node, int needle)
{
    if (node == 0)
        return 0;
    else if (needle < node->elem)
        return search_node(node->lchild, needle);
    else if (needle > node->elem)
        return search_node(node->rchild, needle);
    else
        return node;
}

int
avl_search(avltree_t* tree, int needle)
{
    return search_node(tree->root, needle) != 0;
}

void
graph_nodes(node_t* root, int level)
{
    int i;
    if (root == 0)
        return;
    graph_nodes(root->rchild, level + 1);
    for (i = 0; i < level; ++i)
        printf("    ");
    printf("%d\n", root->elem, root->height);
    graph_nodes(root->lchild, level + 1);
}

void
avl_graph(avltree_t* tree)
{
    printf("-----------------Tree-----------------\n");
    graph_nodes(tree->root, 0);
}

[1] AVL 树的高度已被证明小于 log_ϕ(2)*log₂(n+2)-1，ϕ 是黄金分割，n 是节点（元素）数量。