[Leetcode] 295. Find Median from Data Stream

最新推荐文章于 2025-01-26 04:30:00 发布

原创最新推荐文章于 2025-01-26 04:30:00 发布 · 214 阅读

1 ·

CC 4.0 BY-SA版权

Leetcode 专栏收录该内容

294 篇文章

订阅专栏

本文介绍了一种使用两个优先级队列（一个最大堆和一个最小堆）来维护流数据中位数的有效方法。该算法能够实时更新数据流中的中位数，适用于数据不断变化的场景。

这道题出现在这个挺好其实挺匪夷所思的，这是很早期（大约2012年我找工作的时候）的经典题目之一，当初没有出现在leetcode上我还以为是因为太经典leetcode就放弃了。做法也是非常固定。就是保持两个priorityqueue，一个最大堆（也就是top element是当前堆最大值），一个最小堆（也就是top element是当前堆最小值）。最大堆保持当前stream的数字的较小的一半的部分，最小堆保持当前stream的数字组里较大的一半。这样的话，你最大堆里面top element就是median本身（stream的数字个数为奇数）或者刚好小于median（如果stream的数字个数为偶数）的一个数，最小堆的top elment就是刚好大于median的那个数字。

所以当stream每input一个数字的时候，分为以下几种情况
1. 如果两个堆都是空的，就放到最大堆里（最小堆也可以，那么当总数为奇数的时候，median就在最小堆里，也就是如果最小堆或者最大堆在stream总数为奇数的时候，总有一个堆刚好比另一个堆多出一个数字，那个多出来的那边的top element就是median）
2.否则，如果当前数字大于最小堆的top element，你就明白这个数字其实应该是属于median较大的那一侧，此时你就放到最小堆里，如果这个时候最小堆的size大于最大堆，你就把最小堆的top element放到最大堆里（保持median永远出自最大堆的consistency）
3. 如果当前数字小于等于最大堆的top element，这个数字就属于median较小的那一侧，你就得首先先看一下最大堆的size是不是大于最小堆，是的话就先把最大堆的那个top element放到最小堆里，再push当前数字到最大堆，这是为了保持两边相等或者最小堆最多比最大堆多一个数字（多出来的top element部分就是median）
4. 当取median的时候，如果最大堆比最小堆多一个数字，也就是表示stream当前的input总数为奇数，直接取最大堆的top element，否则就取两者的平均值。

根据上述描述，可以得到算法如下

class MedianFinder {

    /** initialize your data structure here. */
    PriorityQueue<Integer> min, max;
    public MedianFinder() {
        min = new PriorityQueue<Integer>();
        max = new PriorityQueue<Integer>(10, Collections.reverseOrder());
    }
    
    public void addNum(int num) {
        if (min.isEmpty() && max.isEmpty()) {
            max.offer(num);
        } else {
            if (num > max.peek()) {
                min.offer(num);
                if (min.size() > max.size()) {
                    max.offer(min.remove());
                }
            } else {
                if (max.size() > min.size()) {
                    min.offer(max.remove());
                }
                
                max.offer(num);
            }
        }
    }
    
    public double findMedian() {
        if (max.size() > min.size()) {
            return max.peek();
        } else {
            return (max.peek() + min.peek()) / (double)2.0;
        }
    }
}

此外要注意的是，PriorityQueue默认是最小堆，也就是top element是当前堆里的最小值，用了Collections.reverseOrder作为comparator之后才会变成最大堆。