数据流中的中位数(优先队列/堆)

183 阅读3分钟

题目

如何得到一个数据流中的中位数?如果从数据流中读出奇数个数值,那么中位数就是所有数值排序之后位于中间的数值。如果从数据流中读出偶数个数值,那么中位数就是所有数值排序之后中间两个数的平均值。

例如,

[2,3,4] 的中位数是 3

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

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

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

输入:
["MedianFinder","addNum","addNum","findMedian","addNum","findMedian"]
[[],[1],[2],[],[3],[]]
输出:[null,null,null,1.50000,null,2.00000]

思路

给定一长度为 N 的无序数组,其中位数的计算方法:首先对数组执行排序(使用 O(NlogN) 时间),然后返回中间元素即可(使用 O(1) 时间)。

针对本题,根据以上思路,可以将数据流保存在一个列表中,并在添加元素时 保持数组有序 。此方法的时间复杂度为 O(N)O(N) ,其中包括: 查找元素插入位置 O(\log N)O(logN) (二分查找)、向数组某位置插入元素 O(N)O(N) (插入位置之后的元素都需要向后移动一位)。

优化(利用堆)

建立一个 小顶堆 A 和 大顶堆 B ,各保存列表的一半元素,且规定:

  • A 保存 较大 的一半,长度为N/2N/2(N为偶数)或者N+1/2(N+1)/2(N为奇数);
  • B 保存 较小 的一半,长度为N/2N/2(N为偶数)或者N1/2(N-1)/2(N为奇数);

这样的选择的好处在于,小顶堆的堆顶是较大值一半的最小值,也就是中间的值,而大顶堆的堆顶是和偶数一样的中间值,奇数中间值的下一值。

addNum(num) 函数:

在插入顶堆的有两个注意事项:

  1. 当 m=n(即N为偶数时):需要向A添加一个元素。但是由于元素num可能属于B中的元素,因此不可以直接将num插入A中。应该先将num插入B,再将B中的堆顶元素插入至A。这样就可以始终保持A保存较大的一半,B保存较小的一半。
  2. 当 m != n(即N为奇数时):显然这时A中的元素多余B中的元素,应该向B中添加元素,但是有可能添加的元素num属于A中的元素。应该先将num插入A,再将A中的堆顶元素插入至B。这样就可以始终保持A保存较大的一半,B保存较小的一半。

findMedian() 函数:

查找中位数时,我们要根据奇数和偶数返回,如果是偶数,则取A,B的堆顶。如果是奇数,则取A堆顶,因为A比B多1.所以取A堆顶。

代码

class MedianFinder {
    PriorityQueue<Integer> min,max; //定义两个堆,小顶堆
    /** initialize your data structure here. */
    public MedianFinder() {
        min = new PriorityQueue<>();
        max = new PriorityQueue<>((x,y)->(y-x));
    }
    public void addNum(int num) {
        if(min.size() != max.size()){
            min.add(num);
            max.add(min.poll());
        }else{
            max.add(num);
            min.add(max.poll());
        }

    }
    public double findMedian() {
        if(min.size() != max.size()){
            //double s = min.poll();
            //int s1 = (int) s;
            //min.add(s1);
            //return s;
            return min.peek();
        }else{
            double a = min.poll();
            double b = max.poll();
            int a1 = (int) a;
            int b1 = (int) b;
            min.add(a1);
            max.add(b1);
            return (a+b)/2;
            return (min.peek()+max.peek())/2.0;
        }
    }
}