30017：插入型 AVL 树

题目

实现支持插入、查找和完整不变式验证的 AVL 整数树。

解析

每个结点保存子树高度。插入沿搜索路径递归返回；只有子树插入成功后才更新高度并旋转，所以分配失败会沿调用链直接返回，不会执行任何旋转。平衡因子为左子树高度减右子树高度，四种失衡分别由一次单旋或一次双旋恢复。

验证器重新计算每个结点的高度，同时检查严格搜索顺序和左右高度差，不能只检查根结点的平衡。

答案

#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>

typedef struct AvlTree AvlTree;
typedef struct AvlNode AvlNode;

struct AvlNode {
    int64_t key;
    size_t height;
    AvlNode *left;
    AvlNode *right;
};

struct AvlTree {
    AvlNode *root;
    size_t size;
};

static size_t node_height(const AvlNode *node) {
    return node == NULL ? 0 : node->height;
}

static size_t maximum_size(size_t left, size_t right) {
    return left > right ? left : right;
}

static void update_height(AvlNode *node) {
    size_t child_height = maximum_size(
        node_height(node->left), node_height(node->right)
    );
    node->height = child_height == SIZE_MAX
        ? SIZE_MAX
        : child_height + 1;
}

static AvlNode *rotate_right(AvlNode *root) {
    AvlNode *pivot = root->left;
    root->left = pivot->right;
    pivot->right = root;
    update_height(root);
    update_height(pivot);
    return pivot;
}

static AvlNode *rotate_left(AvlNode *root) {
    AvlNode *pivot = root->right;
    root->right = pivot->left;
    pivot->left = root;
    update_height(root);
    update_height(pivot);
    return pivot;
}

static AvlNode *rebalance(AvlNode *node) {
    update_height(node);
    size_t left_height = node_height(node->left);
    size_t right_height = node_height(node->right);

    if (left_height > right_height &&
        left_height - right_height > 1) {
        if (node_height(node->left->right) >
            node_height(node->left->left)) {
            node->left = rotate_left(node->left);
        }
        return rotate_right(node);
    }
    if (right_height > left_height &&
        right_height - left_height > 1) {
        if (node_height(node->right->left) >
            node_height(node->right->right)) {
            node->right = rotate_right(node->right);
        }
        return rotate_left(node);
    }
    return node;
}

static bool insert_node(
    AvlNode **link,
    int64_t key,
    bool *out_inserted
) {
    if (*link == NULL) {
        AvlNode *node = malloc(sizeof *node);
        if (node == NULL) {
            return false;
        }
        node->key = key;
        node->height = 1;
        node->left = NULL;
        node->right = NULL;
        *link = node;
        *out_inserted = true;
        return true;
    }
    if (key == (*link)->key) {
        *out_inserted = false;
        return true;
    }

    AvlNode **next = key < (*link)->key
        ? &(*link)->left
        : &(*link)->right;
    if (!insert_node(next, key, out_inserted)) {
        return false;
    }
    if (*out_inserted) {
        *link = rebalance(*link);
    }
    return true;
}

AvlTree *avl_create(void) {
    AvlTree *tree = malloc(sizeof *tree);
    if (tree == NULL) {
        return NULL;
    }
    tree->root = NULL;
    tree->size = 0;
    return tree;
}

static void destroy_nodes(AvlNode *node) {
    if (node == NULL) {
        return;
    }
    destroy_nodes(node->left);
    destroy_nodes(node->right);
    free(node);
}

void avl_destroy(AvlTree *tree) {
    if (tree == NULL) {
        return;
    }
    destroy_nodes(tree->root);
    free(tree);
}

bool avl_insert(
    AvlTree *tree,
    int64_t key,
    bool *out_inserted
) {
    if (tree == NULL || out_inserted == NULL) {
        return false;
    }
    bool inserted = false;
    if (!insert_node(&tree->root, key, &inserted)) {
        return false;
    }
    if (inserted) {
        ++tree->size;
    }
    *out_inserted = inserted;
    return true;
}

bool avl_contains(const AvlTree *tree, int64_t key) {
    if (tree == NULL) {
        return false;
    }
    const AvlNode *node = tree->root;
    while (node != NULL) {
        if (key == node->key) {
            return true;
        }
        node = key < node->key ? node->left : node->right;
    }
    return false;
}

static bool validate_node(
    const AvlNode *node,
    bool has_lower,
    int64_t lower,
    bool has_upper,
    int64_t upper,
    size_t *out_height,
    size_t *out_count
) {
    if (node == NULL) {
        *out_height = 0;
        *out_count = 0;
        return true;
    }
    if ((has_lower && node->key <= lower) ||
        (has_upper && node->key >= upper)) {
        return false;
    }
    size_t left_height = 0;
    size_t right_height = 0;
    size_t left_count = 0;
    size_t right_count = 0;
    if (!validate_node(
            node->left, has_lower, lower, true, node->key,
            &left_height, &left_count
        ) ||
        !validate_node(
            node->right, true, node->key, has_upper, upper,
            &right_height, &right_count
        )) {
        return false;
    }
    if ((left_height > right_height && left_height - right_height > 1) ||
        (right_height > left_height && right_height - left_height > 1)) {
        return false;
    }
    size_t expected = maximum_size(left_height, right_height);
    if (expected == SIZE_MAX) {
        return false;
    }
    ++expected;
    if (node->height != expected ||
        right_count > SIZE_MAX - left_count ||
        left_count + right_count == SIZE_MAX) {
        return false;
    }
    *out_height = expected;
    *out_count = left_count + right_count + 1;
    return true;
}

bool avl_validate(const AvlTree *tree) {
    if (tree == NULL) {
        return false;
    }
    size_t height = 0;
    size_t count = 0;
    return validate_node(
        tree->root, false, 0, false, 0, &height, &count
    ) && count == tree->size;
}

插入和查找为 $O(\log n)$，验证为 $O(n)$；每次旋转只修改局部链接并重新计算局部高度。