10. AVL树：实现与测试

# 10. AVL树：实现与测试

Redis常被当作键值存储，但Redis里“值”这部分可不限于普通字符串，列表、哈希表（hashmap）和有序集合（sorted set）都超好用。由于Redis拥有丰富的数据结构，它还被称为“数据结构服务器”。

Redis常被用作内存缓存，在内存里存数据时，就可以放开了使用各种数据结构。Redis里的有序集合数据结构那可是相当独特又实用。它不仅能给数据排序，还能按排名查询有序数据。要是往有序集合里存2000万条记录，你不用遍历前1000万条，就能直接拿到排名第1000万的记录，这本事目前的SQL数据库可模仿不来。

顾名思义，“有序集合”是用来排序的。树，尤其是平衡二叉树，是存储有序数据的热门数据结构。在众多数据结构里，作者发现AVL树特别简单，代码写起来也轻松，本书就用它来实现有序集合。真正的Redis项目用的是跳表（skiplist），它写起来也不难。

AVL树的核心思路是限制左右子树的高度差。子树之间的高度差最多只能是1，绝不能达到2。往AVL树里插入或删除节点时，高度差可能会暂时达到2，不过可以通过节点旋转来调整。旋转操作是平衡二叉树的基础，像红黑树（RB tree）这些平衡树也会用到。旋转之后，高度差为2的节点就又变回高度差最多为1了。

咱们先从树节点开始：

struct AVLNode {
    uint32_t depth = 0;
    uint32_t cnt = 0;
    AVLNode *left = NULL;
    AVLNode *right = NULL;
    AVLNode *parent = NULL;
};

static void avl_init(AVLNode *node)  {
    node->depth = 1;
    node->cnt = 1;
    node->left = node->right = node->parent = NULL;
}
1
2
3
4
5
6
7
8
9
10
11
12
13

这就是个普通的二叉树节点，加了几个额外的字段。depth字段表示树的高度，cnt字段表示树的节点数。这个cnt字段不是AVL树特有的，它是用来实现基于排名的查询的，下一章会详细讲。

下面是一些辅助函数：

static uint32_t avl_depth(AVLNode *node)  {
    return node? node->depth : 0;
}

static uint32_t avl_cnt(AVLNode *node)  {
    return node? node->cnt : 0;
}

static uint32_t max(uint32_t lhs, uint32_t rhs)  {
    return lhs < rhs? rhs : lhs;
}

// 维护depth和cnt字段
static void avl_update(AVLNode *node)  {
    node->depth = 1 + max(avl_depth(node->left), avl_depth(node->right));
    node->cnt = 1 + avl_cnt(node->left) + avl_cnt(node->right);
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

接下来是节点旋转的代码：

static AVLNode * rot_left(AVLNode *node)  {
    AVLNode *new_node = node->right;
    if (new_node->left)  {
        new_node->left->parent = node;
    }
    node->right = new_node->left;
    new_node->left = node;
    new_node->parent = node->parent;
    node->parent = new_node;
    avl_update(node);
    avl_update(new_node);
    return new_node;
}

static AVLNode * rot_right(AVLNode *node)  {
    // 是rot_left()的镜像操作
    // 代码省略...
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

左旋（rot_left）操作的可视化展示：

    b           d
   / \         /
  a   d  ==>  b
       \     / \
        c   a   c
1
2
3
4
5

avl_fix_left和avl_fix_right函数是用来修正过度的高度差的：

// 左子树太深
static AVLNode *avl_fix_left(AVLNode * root)  {
    if (avl_depth(root->left->left) < avl_depth(root->left->right))  {
        root->left = rot_left(root->left);
    }
    return rot_right(root);
}

// 右子树太深
static AVLNode *avl_fix_right(AVLNode * root)  {
    if (avl_depth(root->right->right) < avl_depth(root->right->left))  {
        root->right = rot_right(root->right);
    }
    return rot_left(root);
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

要是右子树太深，来个左旋就能解决问题。不过在左旋之前，可能得先对右子树来个右旋，确保右子树“站对方向”。下面是可视化展示：

    b           b           d
   / \         / \         / \
  a   c  ==>  a   d  ==>  b   c
       \     /
        d   a
1
2
3
4
5

avl_fix函数用于在插入或删除操作后修正所有问题。它从最初受影响的节点开始，一直处理到根节点。因为旋转可能会改变树的根节点，所以这个函数会返回根节点。这可是咱们AVL树实现的核心部分。

// 修正不平衡的节点并保持树的性质，直到到达根节点
static AVLNode *avl_fix(AVLNode *node)  {
    while (true)  {
        avl_update(node);
        uint32_t l = avl_depth(node->left);
        uint32_t r = avl_depth(node->right);
        AVLNode **from = NULL;
        if (node->parent)  {
            from = (node->parent->left == node)
                   ? &node->parent->left : &node->parent->right;
        }
        if (l == r + 2)  {
            node = avl_fix_left(node);
        } else if (l + 2 == r)  {
            node = avl_fix_right(node);
        }
        if (!from)  {
            return node;
        }
        *from = node;
        node = node->parent;
    }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

二叉树的插入操作很简单，从根节点开始往下找，找到空的子树就把新节点放进去，然后调用avl_fix函数维护树的平衡。

删除操作就复杂点儿了。要是目标节点没有子树，直接删掉就行；要是只有一个子树，就用这个子树替换掉目标节点。但要是节点有两个子树，就不能直接删了，得删掉它右子树里的“兄弟节点”，然后把这个“兄弟节点”和目标节点交换。下面是删除节点的函数：

// 从树中删除一个节点并返回新的根节点
static AVLNode *avl_del(AVLNode *node)  {
    if (node->right == NULL)  {
        // 没有右子树，用左子树替换该节点
        // 将左子树连接到父节点
        AVLNode *parent = node->parent;
        if (node->left)  {
            node->left->parent = parent;
        }
        if (parent)  {
            // 将左子树连接到父节点
            (parent->left == node? parent->left : parent->right) = node->left;
            return avl_fix(parent);
        } else  {
            // 正在删除根节点？
            return node->left;
        }
    } else  {
        // 用它的下一个兄弟节点交换该节点
        AVLNode *victim = node->right;
        while (victim->left)  {
            victim = victim->left;
        }
        AVLNode * root = avl_del(victim);

        *victim = *node;
        if (victim->left)  {
            victim->left->parent = victim;
        }
        if (victim->right)  {
            victim->right->parent = victim;
        }

        AVLNode *parent = node->parent;
        if (parent)  {
            (parent->left == node? parent->left : parent->right) = victim;
            return root;
        } else  {
            // 正在删除根节点？
            return victim;
        }
    }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43

这是二叉树删除节点的通用函数，里面用到了AVL树特有的avl_fix函数。

熟悉红黑树的小伙伴可能会发现，AVL树的实现代码又少又简单。红黑树删除节点的维护代码可比插入复杂多了，而AVL树插入和删除都用avl_fix函数，这种对称性让写AVL树代码轻松多啦。

AVL树可比咱们之前写的哈希表复杂多了，所以得花更多时间来测试。测试代码也能展示AVL树这些函数怎么用。

下面是测试用的数据类型。要是你对侵入式数据结构不太熟，可以去看看哈希表那一章。

struct Data {
    AVLNode node;
    uint32_t val = 0;
};

struct Container {
    AVLNode * root = NULL;
};
1
2
3
4
5
6
7
8

插入节点的代码：

static void add(Container &c, uint32_t val)  {
    Data *data = new Data();
    avl_init(&data->node);
    data->val = val;

    if (!c.root)  {
        c.root = &data->node;
        return;
    }

    AVLNode *cur = c.root;
    while (true)  {
        AVLNode **from =
                (val < container_of(cur, Data, node)->val)
                ? &cur->left : &cur->right;
        if (!*from)  {
            *from = &data->node;
            data->node.parent = cur;
            c.root = avl_fix(&data->node);
            break;
        }
        cur = *from;
    }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

下面展示删除节点的操作：

static bool del(Container &c, uint32_t val)  {
    AVLNode *cur = c.root;
    while (cur)  {
        uint32_t node_val = container_of(cur, Data, node)->val;
        if (val == node_val)  {
            break;
        }
        cur = val < node_val? cur->left : cur->right;
    }
    if (!cur)  {
        return false;
    }

    c.root = avl_del(cur);
    delete container_of(cur, Data, node);
    return true;
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

下面这个函数用来验证树结构是否正确：

static void avl_verify(AVLNode *parent, AVLNode *node)  {
    if (!node)  {
        return;
    }

    assert(node->parent == parent);
    avl_verify(node, node->left);
    avl_verify(node, node->right);

    assert(node->cnt == 1 + avl_cnt(node->left) + avl_cnt(node->right));

    uint32_t l = avl_depth(node->left);
    uint32_t r = avl_depth(node->right);
    assert(l == r || l + 1 == r || l == r + 1);
    assert(node->depth == 1 + max(l, r));

    uint32_t val = container_of(node, Data, node)->val;
    if (node->left)  {
        assert(node->left->parent == node);
        assert(container_of(node->left, Data, node)->val <= val);
    }
    if (node->right)  {
        assert(node->right->parent == node);
        assert(container_of(node->right, Data, node)->val >= val);
    }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

比较AVL树内容和预期数据的代码：

static void extract(AVLNode *node, std::multiset<uint32_t> &extracted)  {
    if (!node)  {
        return;
    }
    extract(node->left, extracted);
    extracted.insert(container_of(node, Data, node)->val);
    extract(node->right, extracted);
}

static void container_verify(
        Container &c, const std::multiset<uint32_t> &ref)
{
    avl_verify(NULL, c.root);
    assert(avl_cnt(c.root) == ref.size());
    std::multiset<uint32_t> extracted;
    extract(c.root, extracted);
    assert(extracted == ref);
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

测试完可别忘了清理：

static void dispose(Container &c)  {
    while (c.root)  {
        AVLNode *node = c.root;
        c.root = avl_del(c.root);
        delete container_of(node, Data, node);
    }
}
1
2
3
4
5
6
7

咱们的测试用例从简单的开始：

Container c;

// 一些快速测试
container_verify(c, {});
add(c, 123);
container_verify(c, {123});
assert(!del(c, 124));
assert(del(c, 123));
container_verify(c, {});

// 顺序插入
std::multiset<uint32_t> ref;
for (uint32_t i = 0; i < 1000; i += 3)  {
    add(c, i);
    ref.insert(i);
    container_verify(c, ref);
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

然后来点随机操作：

// 随机插入
for (uint32_t i = 0; i < 100; i++)  {
    uint32_t val = (uint32_t)rand() % 1000;
    add(c, val);
    ref.insert(val);
    container_verify(c, ref);
}

// 随机删除
for (uint32_t i = 0; i < 200; i++)  {
    uint32_t val = (uint32_t)rand() % 1000;
    auto it = ref.find(val);
    if (it == ref.end())  {
        assert(!del(c, val));
    } else  {
        assert(del(c, val));
        ref.erase(it);
    }
    container_verify(c, ref);
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

还有些更有针对性的测试。给定一定大小的树，在每个可能的位置进行插入和删除操作。

static void test_insert(uint32_t sz)  {
    for (uint32_t val = 0; val < sz; ++val)  {
        Container c;
        std::multiset<uint32_t> ref;
        for (uint32_t i = 0; i < sz; ++i)  {
            if (i == val)  {
                continue;
            }
            add(c, i);
            ref.insert(i);
        }
        container_verify(c, ref);

        add(c, val);
        ref.insert(val);
        container_verify(c, ref);
        dispose(c);
    }
}

static void test_remove(uint32_t sz)  {
    for (uint32_t val = 0; val < sz; ++val)  {
        Container c;
        std::multiset<uint32_t> ref;
        for (uint32_t i = 0; i < sz; ++i)  {
            add(c, i);
            ref.insert(i);
        }
        container_verify(c, ref);

        assert(del(c, val));
        ref.erase(val);
        container_verify(c, ref);
        dispose(c);
    }
}

// 在不同位置进行插入/删除操作
for (uint32_t i = 0; i < 200; ++i)  {
    test_insert(i);
    test_remove(i);
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42

多亏了这些测试用例，作者在写这章的时候发现并修正了好几个错误。

# 练习：

  1. 虽说咱们的AVL树代码不算多，但这个实现可能效率不太够。代码里有些指针更新操作是多余的，这或许能优化一下。而且，为了保持平衡，其实不用存高度值，存高度差就行。去研究研究高效的AVL树实现吧。
  2. 你能想出更多测试用例吗？这章给出的测试用例可能不太够哦。
  • avl.cpp
  • test_avl.cpp
上次更新: 2025/03/25, 00:48:42
09. 数据序列化
11. AVL树和有序集合

社区交流

讨论与留言

前往 GitHub Issues →

本留言区仅对应当前文章，欢迎补充观点、提出问题或帮助修正文中疏漏。 留言由 GitHub/Gitalk 提供，需要使用 GitHub 登录。

最近更新
第 1 课：专栏导论与安全边界
07-01
第 3 课：安装编译与调试环境
07-01
第 4 课：第三方库下载、编译与依赖管理
07-01
更多文章>

Copyright © 2024-2026 沪ICP备2023015129号 张小方 版权所有