树的用途来自额外约束。二叉搜索树要求左子树键小于节点、右子树键大于节点;最小堆只要求父节点不大于孩子。搜索树可支持有序查询,堆则擅长反复取当前最小值。把最小堆按数组下标直接打印,并不会得到完整有序序列。

搜索树为什么需要平衡

搜索树查询沿根到叶的一条路径,成本与高度有关。依次插入有序键会让普通BST退化为链,高度 O(n)。AVL树通过高度约束和旋转控制高度,红黑树用颜色约束达到同类目标;旋转改变局部连接,但保持中序顺序。因此平衡树的重点是维护约束,而不是记住一种树形图。

范围查询可以先找下界,再按序输出,成本应包含输出数量 k。如果数据很少修改,有序数组加二分可能更简单。B树用更高分支减少外存访问层数;跳表用随机多级索引支持有序查询。它们都要结合数据分布、并发与存储介质选型。

完全二叉堆可以放在数组中

零起始数组下标 i 的孩子是 2*i+1 和 2*i+2。插入时从尾部向上修复,弹出根时把末尾元素搬到根再向下修复,路径长度 O(log n)。自底向上建堆是 O(n):大多数节点靠近叶子,只需很短的修复路径,不能把每个节点都按最坏高度粗算为最终紧确界。

Python heapq 直接对列表维护最小堆,不会封装成另一个容器。只要需要最大值,未必必须依赖新版本API;保留最大的K个数,可以直接使用最小堆。

import heapq


def top_k(values, k):
    if k < 0:
        raise ValueError('k must be non-negative')
    if k == 0:
        return []
    heap = []
    for value in values:
        if len(heap) < k:
            heapq.heappush(heap, value)
        elif value > heap[0]:
            heapq.heapreplace(heap, value)
    return sorted(heap, reverse=True)


if __name__ == '__main__':
    print(top_k(iter([4, 1, 8, 8, 2]), 3))

堆中保存已读元素中最大的至多K个。未满时加入,满后只有大于堆顶的新值才值得替换。替换掉候选中最小值,保持不变量。输入 [4,1,8,8,2]、K=3返回 [8,8,4],重复记录按不同元素保留。

扫描成本 O(n log K),最终降序输出另需 O(K log K)。K=0返回空;负K明确拒绝。K大于输入长度时保留全部输入,函数也接受一次性迭代器,不先把整条流转为列表。

同优先级任务的细节

堆里存 (priority, task) 时,优先级相同会继续比较 task。如果任务对象不可比较,会抛异常;若要求先进先出,可存 (priority, sequence_number, task)。任务优先级变化时不能只改对象字段,必须重新修复堆,或加入新条目并在弹出时跳过失效旧条目。

练习与面试自查

改为统计最常见的K个词,并明确同频时按什么排序。再设计任务取消标记,说明失效条目过多时如何压缩堆。面试中应能解释:堆为何不能 O(log n) 查任意值,平衡树的旋转为何不破坏排序,以及建堆 O(n) 与逐个插入 O(n log n) 的差别。

本地运行

示例仅使用 Python 标准库,适用于 Python 3.10 及以上。把 example.py 与 test_example.py 保存到同一个目录,在该目录运行:

python example.py
python -m unittest -v test_example.py

测试检查本文列出的边界;新增练习仍需自己补充断言。

参考资料