数据流中位数的双堆解法：从超时到O(log n)的工程实践

数据流的中位数LeetCode 295Median of Data Stream在LeetCode热题100乃至各大厂的算法面试里属于那种“看着简单一写就乱”的题目。我第一次做它的时候第一反应是维护一个有序列表每次插入之后取中间下标——测了几个用例确实能过但提交的时候直接超时。那时候我才意识到“数据流”三个字的重量不在“中位数”而在“数据流”数据是源源不断进来的每进来一个数都要能立刻回答当前的中位数这才是这题真正的考点。而双堆法就是专门用来打破“动态插入”和“快速查询”这对矛盾的标准答案。这篇文章我会从暴力解法为什么不行的根因讲起把双堆法的结构设计、Java代码实现、边界条件、复杂度证明以及面试里可能被追问的变体问题一次讲透。1. 题目长什么样为什么排序法在这里“废了”1.1 一句话看懂LeetCode 295题目要求设计一个类支持两个操作void addNum(int num)从数据流中加入一个整数double findMedian()返回当前所有元素的中位数中位数的定义不用多说元素个数为奇数时取排序后中间那个数为偶数时取中间两个数的平均值。举个例子数据流依次进来[1, 2, 3]此时中位数是2继续进来一个4变成[1, 2, 3, 4]中位数就是(2 3) / 2 2.5。这个题在LeetCode热题100里被归类为“堆”类题目题号是295Java版本在面试中出镜率极高。它考察的不是某个偏门API而是你能否根据业务场景选对数据结构——数据流、动态插入、频繁查询中位数这三个关键词放在一起你首先想到的必须是有序性可以被“拆成两半”来维护的结构而不是排序。1.2 直接排序为什么不可行很多人的第一版代码长这样用ArrayList存所有数字addNum直接加到末尾findMedian调用Collections.sort()排个序然后按下标取中间。这种写法在小规模数据下完全没问题但往LeetCode的极限测试数据上一跑就原形毕露。Collections.sort()的时间复杂度是O(n log n)如果addNum和findMedian交替调用n次总代价会膨胀到O(n^2 log n)数据量到几万级别基本就跑不动了。还有人会想出更“优化”的版本每次addNum都保持列表有序用二分查找找到新数的插入位置然后用System.arraycopy把后面的元素整体右移。这个方案的findMedian确实能到O(1)但addNum仍然是O(n)因为数组/列表的插入天生要搬移元素。为什么排序法在普通数组里没问题到数据流里就不行了关键就在“动态”两个字。静态数据你排一次序就完事了数据流是边进边问每一次插入都可能紧跟一次查询插入代价必须压低到O(log n)才能扛住规模。这其实是在线算法和离线算法的区别。离线时你可以先拿到全量数据排好序再慢慢查在线时必须每来一个数就立刻维护好“未来可能被查询到的信息”。中位数这种统计量在在线场景下的经典解法就是双堆法。2. 双堆法把“中位数”翻译成两个堆的数据结构2.1 最大堆和最小堆的分工双堆法的核心思想用一句话概括把当前所有元素分成“较小的一半”和“较大的一半”较小的一半装进一个最大堆maxHeap较大的一半装进一个最小堆minHeap。maxHeap的堆顶是“较小一半”里面的最大值也就是所有小于等于中位数的数里最大的那个minHeap的堆顶是“较大一半”里面的最小值也就是所有大于等于中位数的数里最小的那个当两个堆的元素数量相等时中位数就是两个堆顶的平均值当maxHeap比minHeap多一个元素时中位数就是maxHeap的堆顶。这个过程很像一个天平左边放小数右边放大数天平尽量保持平衡而归零或者只偏一格的位置就是整个数据流的中位数所在。2.2 不变量两个堆的大小关系要实现上述效果必须始终维护两条不变量maxHeap.size()与minHeap.size()之差不超过1且这里约定maxHeap.size() minHeap.size()maxHeap中任意一个元素都小于等于minHeap中任意一个元素只要这两条不变量成立中位数的计算就是O(1)的堆顶读取。任何一次addNum之后的堆调整本质都是恢复这两条不变量。很多刚开始学双堆法的人容易有一个误区以为只要把数字随便塞进两个堆再调整一下大小关系就行。这是错误的。如果不变量2被破坏——比如maxHeap里混进一个很大的数minHeap里混进一个很小的数——那么两个堆的大小虽然平衡但堆顶已经不能代表“中间”了算出来的中位数是错的。2.3 为什么这种划分能让中位数“自动”出现在堆顶双堆法的巧妙之处在于它不需要像插入排序那样精确维护全序只需要维护一个“中间分界”。这个分界由两个堆顶构成maxHeap的堆顶是界线上边的最后一个元素minHeap的堆顶是界线下边的第一个元素。打个比方想象一排按身高站队的人中间画一条线。左边的队伍里最高的人站在线边右边的队伍里最矮的人站在线边。无论新来的人多高多矮先让他去左边队伍里站好然后把左边队伍里现在最高的那个人送到右边去——这个“最高的人”就是原先左边队伍的界边。重复这个过程分界线始终正确。这个“先塞左边、再从左边把最大值抛到右边”的操作本质上是一种代价只有O(log n)的两段式插入排序既保持了全体的有序性语义又避开了数组插入时的搬移开销。3. Java实现从PriorityQueue到完整题解3.1 初始化两个堆的创建Java里堆就是PriorityQueue默认是小顶堆也就是堆顶永远是最小元素。要得到大顶堆需要自定义比较器。class MedianFinder { private PriorityQueueInteger maxHeap; private PriorityQueueInteger minHeap; public MedianFinder() { // 大顶堆存较小的一半堆顶是这一半的最大值 maxHeap new PriorityQueue((a, b) - Integer.compare(b, a)); // 小顶堆存较大的一半堆顶是这一半的最小值 minHeap new PriorityQueue(); } }这里有个很多Java新手会犯的错大顶堆比较器直接写(a, b) - b - a。两个int相减在正常数据下没问题但遇到Integer.MAX_VALUE和Integer.MIN_VALUE这类极端值减法会溢出比较结果直接反转堆就坏了。所以要么用Integer.compare(b, a)要么用Comparator.reverseOrder()总之别自己用减法当比较器。3.2 addNum()先塞后洗再平衡public void addNum(int num) { maxHeap.offer(num); minHeap.offer(maxHeap.poll()); if (minHeap.size() maxHeap.size()) { maxHeap.offer(minHeap.poll()); } }这段代码只有三行但每一步都有明确职责maxHeap.offer(num)先把新数放进最大堆minHeap.offer(maxHeap.poll())从maxHeap弹出最大值放进minHeap如果minHeap比maxHeap多了就把minHeap的最小值弹回maxHeap第一眼看上去很多人会困惑为什么要绕这么一圈而不是直接判断num和两个堆顶的大小关系再决定放哪边其实两种思路都能实现但“先塞后洗”这个套路有一个巨大优势它对任何输入都无脑正确不需要处理maxHeap为空、num恰好等于堆顶等边界情况。第二步把maxHeap当前最大值交给minHeap保证了“新来的数无论多大都会先进入较小半区参与比较然后比较大半区的门槛是否合格”。这个过程是自动维持不变量2的。第三步再平衡大小保证不变量1。我实际用下来这个三行版本比“判断左右再插入”的版本要好写很多面试时也不容易漏分支。3.3 findMedian()根据大小关系直接取堆顶public double findMedian() { if (maxHeap.size() minHeap.size()) { return maxHeap.peek(); } return (maxHeap.peek() minHeap.peek()) / 2.0; }奇数个数时maxHeap比minHeap多一个直接返回maxHeap的堆顶。偶数个数时两个堆一样多返回两个堆顶的平均值。注意这里的2.0不能写成2。如果写成整数2Java会做整数除法两个int的和小数部分会被直接丢弃比如(3 4) / 2 3完全错误。这个坑看着小但在面试白板上写代码时特别容易犯。3.4 完整代码import java.util.Comparator; import java.util.PriorityQueue; class MedianFinder { private PriorityQueueInteger maxHeap; private PriorityQueueInteger minHeap; public MedianFinder() { maxHeap new PriorityQueue(Comparator.reverseOrder()); minHeap new PriorityQueue(); } public void addNum(int num) { maxHeap.offer(num); minHeap.offer(maxHeap.poll()); if (minHeap.size() maxHeap.size()) { maxHeap.offer(minHeap.poll()); } } public double findMedian() { if (maxHeap.size() minHeap.size()) { return maxHeap.peek(); } return (maxHeap.peek() minHeap.peek()) / 2.0; } }这份代码直接在LeetCode 295提交就能跑通。Comparator.reverseOrder()返回的是内建的反序比较器比手写 lambda 更干净也不容易出错。4. 边界条件、性能细节与常见坑4.1 大顶堆比较器别用减法这篇文章里我已经第二次强调比较器了因为它真的是PriorityQueue题目里最高频的翻车点。减法b - a在b Integer.MIN_VALUE、a Integer.MAX_VALUE时会溢出为负数比较方向反转堆的内部结构彻底乱掉。实际场景里要触发这个bug需要堆里同时出现极端数值概率并不高但它属于“一旦触发就是灾难级”的bug而且非常难以排查。既然Integer.compare和Comparator.reverseOrder()写起来几乎不费事就直接养成习惯。4.2 处理重复元素数据流里完全可能出现大量重复值比如连续加10个5。双堆法对重复元素天然健壮因为PriorityQueue不要求元素唯一。重复值的堆顶依然是正确的。真正需要注意的是比较器如果实现得不好比如上面说的减法重复元素可能导致堆内部结构在极端情况下混乱。把比较器写正确重复值就不会带来任何额外问题。4.3 处理负数很多人初学者会担心“堆是不是只适合正数”。实际上堆比较的是对象的自然顺序负数的自然顺序本来就在正数前面堆顶的“最小”或“最大”完全按数值大小走负数和正数一样处理不需要任何特殊分支。4.4 空流和初次插入LeetCode的测试数据保证了不会在没有任何元素时调用findMedian但如果你把这段代码拿去二次封装建议自己判空。否则maxHeap.peek()返回null自动拆箱成int会直接抛NullPointerException。另外第一次addNum之后整个流程走一遍maxHeap有一个元素minHeap为空。此时findMedian走“这个分支返回maxHeap.peek()结果正确不需要额外处理首元素。4.5 addNum次数与堆操作的对应关系三行代码里offer和poll各出现了几次最坏情况下第二步的poll和offer各一次第三步如果触发则再各一次所以一次addNum最多执行三次堆操作O(log n)。因为每个操作都是O(log n)所以addNum整体是O(log n)。有人问能不能优化到O(1)不行。因为你要在动态集合里维护有序分界至少需要一次对数的调整。这在理论上也是下界。5. 复杂度分析、正确性证明与面试追问5.1 时间与空间复杂度方案addNumfindMedian每次全量排序O(1)O(n log n)插入排序保持有序O(n)O(1)双堆法O(log n)O(1)数据流场景下addNum和findMedian的调用次数往往是同一个量级。如果各有n次调用全量排序的总代价O(n^2 log n)完全不可接受插入排序的总代价O(n^2)数据量大时也扛不住双堆法的总代价O(n log n)可以平稳跑过5万次操作空间上两个堆一共存储了全部n个元素所以是O(n)。这个空间开销是必须的因为你没有任何理由丢弃历史数据——中位数可能在任何时刻查询任何元素都可能成为中位数。5.2 正确性证明不变量与归纳面试官可能不会直接说“你证明一下”但会问“你怎么保证这个算法是对的”。这时候讲不变量方法最清晰初始化两个堆都为空不变量1和2自然成立添加元素假设添加前两条不变量成立。第一步把num放进maxHeap不变量2可能被破坏吗不会因为即使num非常大它也只在maxHeap里待了一瞬间第二步就会被poll交给minHeap第二步把maxHeap当前最大值交给minHeap。由于maxHeap原本所有元素都小于等于minHeap原本所有元素而新弹出的元素是maxHeap的最大值所以这个值一定小于等于minHeap的所有元素。放进minHeap后两个堆之间的顺序性依然成立。注意这里的关键是“新弹出的元素一定介于两个堆的分界处”第三步恢复大小平衡后不变量1成立用归纳法可以严格证明每一步都维持两条不变量。面试中不需要写完整的数学证明但要把“为什么第二步能保证顺序性”讲清楚这一般就够用了。5.3 面试官会怎么追问追问通常围绕这几种方向展开“如果数据流中有大量重复值会退化吗”——不会理由见4.2“如果几个堆中某个堆为空怎么办”——maxHeap一定非空因为一旦minHeap超过maxHeap就会被弹回这是流程自动保证的“内存不够存下全部数据怎么办”——这其实已经不是双堆法的应用场景了要换分桶法或者采样法“如果改成求数据流的四分位数、百分位数呢”——可以用多个堆或有序结构思路是“用多个分界点把数据切成多段”本质还是堆顶夹出的分界线模型面试官真正想考察的是你为什么用堆而不是背代码。我建议你在准备时用一句话把核心逻辑讲出来用两个堆分别维护数据流中较小的一半和较大的一半堆顶恰好夹出了中位数所在的分界插入时通过固定流程保证分界永远正确查询时只需要看堆顶。6. 双堆法的变体应用滑动窗口中位数与数据流第K大6.1 LeetCode 480滑动窗口中位数的核心思路LeetCode 295搞定之后下一个很自然的进阶题是480滑动窗口中位数。这道题要求在固定长度的滑动窗口内持续返回中位数。暴力解法是每次窗口滑动都排序复杂度O(n*k log k)大数据量下必然超时。双堆法可以解决但有一个关键难点窗口左端滑出的元素要从对应堆里删除。而PriorityQueue的remove(Object)是O(n)的直接删除会拖慢整体。标准解法是懒删除lazy deletion要删除的元素先不真删而是记录一个待删标记等它出现在堆顶时再弹出。配合两个堆之间的大小平衡每次滑动窗口的操作代价仍可做到O(log k)。这个技巧可以理解为一个“延迟支付”的思路删除的成本被推迟到“这个元素恰好挡住答案”的时候才支付而且全流程中每个元素最多被真正删除一次摊还下来依然是O(log k)。如果你想在LeetCode 295之外再练一道堆的实战题480是很合适的选择。6.2 LeetCode 703数据流中的第K大元素如果说中位数要找的是“中间那个数”那第K大要找的是“从大到小数的第K个”。核心思路非常类似用一个小顶堆维护当前最大的K个数堆顶就是第K大。当新元素大于堆顶时把堆顶替换掉并调整。这个过程和双堆法里minHeap的维护逻辑几乎一模一样。区别在于这里只需要一个堆因为你要找的“分界点”只在一端有。这进一步说明了双堆法背后的统一模型用堆顶维护动态数据流中的某个分位点中位数是“两侧都有分界”的特例第K大是“只在一侧有分界”的特例。6.3 双堆法还能干什么“维护动态数据流的两端极值”这个能力让双堆法不只是用来算中位数求数据流中所有元素的中位数本题求数据流中的前K小和前K大配合懒删除求滑动窗口的百分位数求动态集合的中间偏大值、中间偏小值等变体掌握双堆法的关键是掌握“一个有序的中间分界可以被两个堆堆顶夹出来”这件事。无论面试题怎么变最后基本都是这个模型。我自己在实际刷题中的体会是LeetCode 295的代码很短但你如果只看代码而不理解“先塞后洗”的分界维护机制遇到480这种需要变通的题目时会很吃力。所以强烈建议你把本文第二节的不变量分析和第五节正确性证明完整过一遍这样不管题目在哪一层换皮你都能直接拆出它的核心结构。