小木的算法日记-亲手揭秘 TreeMap/TreeSet 背后的魔法 —— 二叉搜索树

最新推荐文章于 2025-06-12 20:27:43 发布

木旭林晖

最新推荐文章于 2025-06-12 20:27:43 发布

阅读量1.2k

点赞数 49

文章标签：算法

本文链接：https://blog.csdn.net/GetOver_/article/details/148512192

版权

告别无序：亲手揭秘 TreeMap/TreeSet 背后的魔法 —— 二叉搜索树

你好，我是小木。在上一篇《几种常见的二叉树类型》中，我们认识了二叉搜索树（BST）。今天，我们将更进一步，揭开 Java 中 TreeMap 和 TreeSet 的神秘面纱。

你是否曾想过：HashMap（哈希映射-类似于字典，字典的时间复杂度是O（1））已经那么快了，为什么我们还需要 TreeMap？

答案是：秩序。HashMap 提供了闪电般的 O(1) 访问速度，但它的内部是无序的。而 TreeMap 不仅能高效存取，还能始终保持键的有序状态。这赋予了它许多 HashMap 无法企及的强大能力。

🌟 二叉搜索树（BST）：TreeMap 的核心引擎

TreeMap 的所有魔法都源于其底层数据结构——二叉搜索树（BST）。

BST 有一个极其重要的特性，我称之为“左小右大”原则：

对于树中的任意一个节点，其左子树的所有节点值都小于它，右子树的所有节点值都大于它。

看下面这棵 BST，是不是一目了然？

这个特性就是 BST 的“超能力”。当我们需要查找数据时，不再需要像无头苍蝇一样遍历所有节点。每一次比较，我们都能排除掉一半的可能！

想象一下，在普通二叉树（O(N)）和 BST（O(logN)）中查找元素 5 的区别：

普通二叉树：得一个一个问：“你是 5 吗？”，直到问遍所有节点。
BST：
1. 从根节点 7 开始，5 < 7，果断往左走。
2. 来到节点 4，5 > 4，果断往右走。
3. 成功找到 5！

这种“分而治之”的查找方式，效率极高。而增、删、改操作也都依赖于查找，因此它们的理想时间复杂度都是 O(logN)。

🛠️ 从 TreeNode 到 TreeMap：构建有序的键值世界

标准的 TreeNode 结构只关心一个值 val：

      class TreeNode:
    def __init__(self, x: int):
        self.val = x
        self.left = None
        self.right = None

为了实现 TreeMap，我们只需稍作改造，让每个节点能同时存储 key 和 value：

      # K 和 V 是泛型类型，代表 Key 和 Value
class TreeNode:
    def __init__(self, key: K, value: V):
        self.key = key
        self.value = value
        self.left = None
        self.right = None

有了这个结构，我们就能搭建一个功能强大的 MyTreeMap。它不仅包含标准 Map 的功能，还提供了一系列与排序相关的“独门绝技”。

MyTreeMap 的 API 蓝图

      # 这是一个接口蓝图，我们将在后续文章中实现它
class MyTreeMap:

    # --- Map 核心功能 (O(logN)) ---
    def put(self, key: K, value: V) -> None: pass
    def get(self, key: K) -> V: pass
    def remove(self, key: K) -> None: pass
    def contains_key(self, key: K) -> bool: pass

    # --- 有序集合的独特魅力 ---
    def keys(self) -> list[K]: pass           # O(N)，返回有序的键列表
    def first_key(self) -> K: pass            # O(logN)，返回最小键
    def last_key(self) -> K: pass             # O(logN)，返回最大键
    def floor_key(self, key: K) -> K: pass    # O(logN)，找<=key的最大键
    def ceiling_key(self, key: K) -> K: pass  # O(logN)，找>=key的最小键
    
    # --- 更高级的排名与选择 (O(logN)) ---
    def select_key(self, k: int) -> K: pass   # 找排名第 k 的键
    def rank(self, key: K) -> int: pass       # 找键 key 的排名

看到这些 API 是不是已经感觉它很强大了？LinkedHashMap 只能按插入顺序排序，而 TreeMap 则是真正地按键的大小来排序。

🚀 方法原理解析：深入 TreeMap 的“内功心法”

1. 基础操作 (get, put, remove)

这些操作都基于 BST 的“左小右大”原则进行查找，找到目标位置后执行相应的操作。时间复杂度都是 O(logN)。

2. 极值查找 (first_key, last_key)

这非常简单，充分利用了 BST 的特性：

first_key() (最小键)：从根节点出发，一直向左走，走到底就是了。
last_key() (最大键)：从根节点出发，一直向右走，走到底就是了。

对于下面这棵树，最小键是 1，最大键是 12。

3. 有序遍历 (keys)

如何得到一个有序的键列表？答案是 中序遍历。BST 的中序遍历结果天然就是一个有序序列。

中序遍历：先遍历左子树，再访问根节点，最后遍历右子树。

对上面的树进行中序遍历，得到的结果就是 [1, 2, 4, 5, 7, 9, 11, 12]。

4. 范围查找 (floor_key, ceiling_key)

这两个方法就像是 get 的智能版。当精确找不到 key 时，它们不会返回 None，而是返回“最接近”的那个：

floor_key(key)：找到小于或等于 key 的那个最大的键。
ceiling_key(key)：找到大于或等于 key 的那个最小的键。

5. 排名与选择 (select_key, rank) - ✨精髓所在✨

这两个方法是 TreeMap 的一大亮点，但如何高效实现呢？

select_key(k)：查找排名为 k 的键。
rank(key)：查找 key 的排名。

一个朴素的想法是中序遍历，遍历到第 k 个元素就行。但这样复杂度是 O(k)，不够高效。

巧妙的解决方案：给 TreeNode 增加一个 size 字段，记录以当前节点为根的子树有多少个节点。

      class TreeNode:
    def __init__(self, key: K, value: V):
        self.key = key
        self.value = value
        self.size = 1  # 包含自己在内，初始为1
        self.left = None
        self.right = None

有了 size，奇迹发生了：

select_key(k)：
假设我们在根节点 7，它的左子树 size 是 3。这意味着在以 7 为根的这棵树里，有 3 个元素比 7 小。所以，7 自己的排名就是 3 + 1 = 4。
- 如果要找排名第 3 的 (k=3)，3 < 4，说明目标在左子树，我们去左子树找排名第 3 的元素。
- 如果要找排名第 5 的 (k=5)，5 > 4，说明目标在右子树，我们去右子树找排名 5 - 4 = 1 的元素。
每一次选择，都把问题规模缩小一半，实现了 O(logN) 的查找！
rank(key)：
size 同样能加速 rank 的计算。想知道 9 的排名？
1. 从根节点 7 开始，9 > 7。
2. 7 的排名是左子树大小 + 1，即 3 + 1 = 4。
3. 我们去右子树中找 9 的局部排名，再加上 7 的排名 4。
4. 最终 rank(9) = 4 (7的排名) + rank_in_subtree(9)。

虽然这会增加维护成本（每次增删都要更新 size），但为了换取 O(logN) 的极致性能，完全值得！

⚠️ 树的平衡

我们一直说复杂度是 O(logN)，但这有一个重要的前提：树是“平衡”的。

一棵平衡的 BST，左右子树高度大致相当，看起来很“茂盛”：

但如果插入的数据是 1, 2, 3, 4, 5 这样有序的，BST 就会退化成一条“链表”：

在这种极端情况下，树的高度为 N，所有操作的复杂度都退化成了 O(N)！

如何解决？

答案是使用自平衡二叉搜索树，例如大名鼎鼎的 红黑树（Red-Black Tree）。它通过在插入和删除后进行一系列“旋转”操作，始终保持树的近似平衡，从而确保性能稳定在 O(logN)。

红黑树的实现非常复杂，但正是这种复杂性，才铸就了 TreeMap 的稳定与强大。

总结与展望

今天，我们一起揭开了 TreeMap 的神秘面纱：

核心引擎：TreeMap 的高效和有序性源于其底层的二叉搜索树（BST）。
性能关键：BST 的性能取决于其平衡性。理想情况下是 O(logN)，最坏情况下是 O(N)。
工业级方案：实际的 TreeMap 使用红黑树等自平衡结构来保证性能。
独门绝技：通过维护 size 字段，可以高效实现 rank 和 select 等高级功能。

工作中常用 TreeMap/TreeSet 功能星级排序

★★★★★ (核心必会，频繁使用)

有序遍历 (keys/迭代器)
- 一句话说明：这是你选择 TreeMap 而非 HashMap 的首要原因。用于生成有序报告、按时间线处理任务、展示排行榜等。
基本增删查改 (put, get, remove, containsKey)
- 一句话说明：数据容器的基础，和 HashMap 用法一致，但能保证操作后集合依然有序。