题目描述

中位数是有序列表中间的数。如果列表长度是偶数,中位数则是中间两个数的平均值。

例如,[2,3,4] 的中位数是 3;[2,3] 的中位数是 (2 + 3) / 2 = 2.5

设计一个支持以下两种操作的数据结构:

void addNum(int num) - 从数据流中添加一个整数到数据结构中。
double findMedian() - 返回目前所有元素的中位数。


示例:

addNum(1)
addNum(2)
findMedian() -> 1.5
addNum(3) 
findMedian() -> 2

来源:力扣(LeetCode)链接:https://leetcode-cn.com/problems/find-median-from-data-stream

解题思路

我们必然需要有序数据结构,本题的核心思路是使用两个优先级队列

中位数是有序数组最中间的元素算出来的对吧,我们可以把「有序数组」抽象成一个倒三角形,宽度可以视为元素的大小,那么这个倒三角的中部就是计算中位数的元素对吧:

然后我把这个大的倒三角形从正中间切成两半,变成一个小倒三角和一个梯形,这个小倒三角形相当于一个从小到大的有序数组,这个梯形相当于一个从大到小的有序数组。

中位数就可以通过小倒三角和梯形顶部的元素算出来对吧?嗯,你联想到什么了没有?它们能不能用优先级队列表示?小倒三角不就是个大顶堆嘛,梯形不就是个小顶堆嘛,中位数可以通过它们的堆顶元素算出来

梯形虽然是小顶堆,但其中的元素是较大的,我们称其为large,倒三角虽然是大顶堆,但是其中元素较小,我们称其为small

当然,这两个堆需要算法逻辑正确维护,才能保证堆顶元素是可以算出正确的中位数,我们很容易看出来,两个堆中的元素之差不能超过 1

因为我们要求中位数嘛,假设元素总数是n,如果n是偶数,我们希望两个堆的元素个数是一样的,这样把两个堆的堆顶元素拿出来求个平均数就是中位数;如果n是奇数,那么我们希望两个堆的元素个数分别是n/2 + 1n/2,这样元素多的那个堆的堆顶元素就是中位数。

想要往large里添加元素,不能直接添加,而是要先往small里添加,然后再把small的堆顶元素加到large中;向small中添加元素同理

为什么呢,稍加思考可以想明白,假设我们准备向large中插入元素:

如果插入的num小于small的堆顶元素,那么num就会留在small堆里,为了保证两个堆的元素数量之差不大于 1,作为交换,把small堆顶部的元素再插到large堆里。

如果插入的num大于small的堆顶元素,那么num就会成为samll的堆顶元素,最后还是会被插入large堆中。

反之,向small中插入元素是一个道理,这样就巧妙地保证了large堆整体大于small堆,且两个堆的元素之差不超过 1,那么中位数就可以通过两个堆的堆顶元素快速计算了。

思路参考:https://labuladong.gitbook.io/algo/shu-ju-jie-gou-xi-lie/shou-ba-shou-she-ji-shu-ju-jie-gou/shu-ju-liu-zhong-wei-shu

解题代码

class MedianFinder {

    private PriorityQueue<Integer> large;
    private PriorityQueue<Integer> small;

    /** initialize your data structure here. */
    public MedianFinder() {
        // 小顶堆
        large = new PriorityQueue<>();
        // 大顶堆
        small = new PriorityQueue<>((a, b)->{
            return b - a;
        });
    }
    
    public void addNum(int num) {
        if (small.size() >= large.size()) {
            small.offer(num);
            large.offer(small.poll());
        } else {
            large.offer(num);
            small.offer(large.poll());
        }
    }
    
    public double findMedian() {
        if (small.size() > large.size()) {
            return small.peek();
        } else if (small.size() < large.size()) {
            return large.peek();
        } else {
            return (small.peek() + large.peek()) / 2.0;
        }
    }
}

/**
 * Your MedianFinder object will be instantiated and called as such:
 * MedianFinder obj = new MedianFinder();
 * obj.addNum(num);
 * double param_2 = obj.findMedian();
 */