[leetcode]对顶堆，对数时间添加元素，常数时间取中位数(或者第K大的数)

你是否遇到过这样的问题：一个数字源源不断地到来，形成一个数据流，你需要随时能够快速地找到这些数字的中位数？

如果用最朴素的方法，每次查找中位数时都对整个数组进行排序，那么每次操作的时间复杂度将是 O(N log N)。当数据量巨大时，这显然是无法接受的。有没有更高效的方法呢？

答案是肯定的！今天，我们就来介绍一种优雅而高效的数据结构——对顶堆 (Two Heaps)，它能以 O(log N) 的时间复杂度完成添加元素，并以 O(1) 的时间复杂度查询中位数。

什么是对顶堆？

“对顶堆”并不是一种全新的、独立的数据结构，而是巧妙地利用两个堆组合而成的一种结构。具体来说，它由以下两部分组成：

1. 一个大顶堆 (Max-Heap)：用来存储数据流中较小的一半元素。堆顶是这半部分元素中的最大值。
2. 一个小顶堆 (Min-Heap)：用来存储数据流中较大的一半元素。堆顶是这半部分元素中的最小值。

通过这种结构，我们将整个有序的数据集合从中间“切开”，切分的“刀刃”就是中位数。大顶堆维护着左半部分，小顶堆维护着右半部分。

对顶堆如何工作？

为了让这个结构能够正确地找到中位数，我们必须维护两个核心的不变量（Invariants）：

1. 值的平衡：大顶堆中的所有元素都必须小于或等于小顶堆中的所有元素。也就是说 max_heap.top() <= min_heap.top()。
2. 数量的平衡：两个堆的大小（元素数量）要么相等，要么相差不超过 1。

只要我们能在每次添加新元素后都保持这两个不变量，中位数就可以轻松地从两个堆的堆顶得出。

添加一个新元素 (addNum)

添加新元素的步骤是整个结构的核心。假设我们要添加一个新数字 num：

1. 决定插入哪个堆：

   • 如果大顶堆为空，或者 num 小于等于大顶堆的堆顶元素，则将 num 插入大顶堆。
   • 否则，将 num 插入小顶堆。

2. 进行平衡调整 (Rebalance)：插入新元素后，可能会打破“数量平衡”的规则。我们需要进行调整：

   • 如果大顶堆的大小比小顶堆大超过 1（即 max_heap.size() > min_heap.size() + 1），则将大顶堆的堆顶元素弹出，并压入小顶堆。
   • 如果小顶堆的大小比大顶堆大（即 min_heap.size() > max_heap.size()），则将小顶堆的堆顶元素弹出，并压入大顶堆。
   • 注：这里的平衡策略可以有多种，核心是维持数量差不大于1。例如，可以让大顶堆的数量总是等于或比小顶堆多一，这样奇数情况下的中位数总在大顶堆。

经过这两步，值的平衡和数量的平衡都得到了保证。

查找中位数 (findMedian)

由于我们已经维护好了堆的结构，查找中位数变得极其简单，只需要 O(1) 时间：

• 如果元素总数为奇数：那么其中一个堆会比另一个多一个元素。中位数就是那个较多元素的堆的堆顶。根据我们的平衡策略，通常是大顶堆的堆顶。
• 如果元素总数为偶数：两个堆的元素数量相等。中位数就是两个堆顶元素的平均值，即 (max_heap.top() + min_heap.top()) / 2.0。

对顶堆的应用场景

对顶堆最经典的应用就是解决 “数据流中的中位数” 问题（例如 LeetCode 295. Find Median from Data Stream）。

除此之外，它还可以用于解决 “滑动窗口中位数”（LeetCode 480. Sliding Window Median）。虽然这个问题更复杂，因为它涉及到元素的删除操作（堆本身不支持高效删除任意元素，需要通过懒删除等技巧实现），但其核心思想仍然是利用对顶堆来划分数据。

总的来说，任何需要将一个动态集合分成两部分，并频繁查询分界点相关值的问题，都可以考虑使用对顶堆。

代码实现

下面我们分别用 C++ 和 Python 来实现这个“中位数查找器”。

C++ 实现

在 C++ 中，std::priority_queue 是实现堆的利器。默认情况下，它是一个大顶堆。要创建小顶堆，我们需要提供一个额外的比较器 std::greater<T>。

#include <iostream>
#include <vector>
#include <queue>
#include <functional>class MedianFinder {
private:// 大顶堆，存储较小的一半数字std::priority_queue<int> max_heap; // 小顶堆，存储较大的一半数字std::priority_queue<int, std::vector<int>, std::greater<int>> min_heap;public:MedianFinder() {}void addNum(int num) {// 优先放入大顶堆max_heap.push(num);// 为了维持值的平衡，如果大顶堆的顶（最大值）大于小顶堆的顶（最小值）// 则将大顶堆的顶移动到小顶堆if (!max_heap.empty() && !min_heap.empty() && max_heap.top() > min_heap.top()) {int val = max_heap.top();max_heap.pop();min_heap.push(val);}// 维持数量的平衡// 我们让大顶堆的大小总是等于或比小顶堆多1if (max_heap.size() > min_heap.size() + 1) {int val = max_heap.top();max_heap.pop();min_heap.push(val);} else if (min_heap.size() > max_heap.size()) {int val = min_heap.top();min_heap.pop();max_heap.push(val);}}double findMedian() {if (max_heap.size() > min_heap.size()) {// 如果总数是奇数，中位数在大顶堆的堆顶return max_heap.top();} else if (!max_heap.empty()) { // 避免两个堆都为空的情况// 如果总数是偶数，中位数是两个堆顶的平均值return (max_heap.top() + min_heap.top()) / 2.0;} else {return 0.0; // 或者抛出异常}}
};int main() {MedianFinder mf;mf.addNum(1);mf.addNum(2);std::cout << "Median is: " << mf.findMedian() << std::endl; // -> 1.5mf.addNum(3);std::cout << "Median is: " << mf.findMedian() << std::endl; // -> 2mf.addNum(0);std::cout << "Median is: " << mf.findMedian() << std::endl; // -> 1.5return 0;
}

Python 实现

Python 的 heapq 模块只提供了小顶堆。那么如何模拟一个大顶堆呢？一个非常巧妙的技巧是：存入一个数的相反数。这样，最小的相反数就对应着最大的原数。

import heapqclass MedianFinder:def __init__(self):# 大顶堆，存储较小的一半数字（通过存储相反数实现）self.max_heap = []  # 小顶堆，存储较大的一半数字self.min_heap = []def addNum(self, num: int) -> None:# 巧妙的平衡策略：# 1. 先将新元素推入大顶堆# 2. 再将大顶堆的最大元素（堆顶）弹出，推入小顶堆# 3. 如果两个堆的大小不平衡，再从小顶堆中弹出最小元素，推回大顶堆# 为了维持大顶堆特性，存入相反数heapq.heappush(self.max_heap, -num)# 将大顶堆中“最大”的元素（即-num中的最小值）移动到小顶堆val = heapq.heappop(self.max_heap)heapq.heappush(self.min_heap, -val)# 如果小顶堆元素过多，则移动回大顶堆，以保持大顶堆数量 >= 小顶堆数量if len(self.min_heap) > len(self.max_heap):val = heapq.heappop(self.min_heap)heapq.heappush(self.max_heap, -val)def findMedian(self) -> float:# 如果大顶堆元素更多，说明总数是奇数，中位数在大顶堆顶if len(self.max_heap) > len(self.min_heap):# 取出时要取回相反数return -self.max_heap[0]else:# 如果数量相等，说明总数是偶数，中位数是两个堆顶的平均值# 注意 max_heap[0] 需要取反return (-self.max_heap[0] + self.min_heap[0]) / 2.0# 测试
if __name__ == "__main__":mf = MedianFinder()mf.addNum(1)mf.addNum(2)print(f"Median is: {mf.findMedian()}")  # -> 1.5mf.addNum(3)print(f"Median is: {mf.findMedian()}")  # -> 2.0mf.addNum(0)print(f"Median is: {mf.findMedian()}")  # -> 1.5

总结

对顶堆是一个看似简单但功能强大的数据结构。通过将一个大顶堆和一个小顶堆组合起来，它巧妙地将数据流一分为二，并始终将中位数维持在两个堆的“交界处”。这使得它能够在对数时间内添加元素，并在常数时间内查询中位数，是处理动态数据集合问题的绝佳工具。