数据流中的中位数

问题陈述

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

例如，

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

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

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

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

示例1：

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

思路分析

可将数据流保存在一个列表中，并将其有序排列，在添加元素时保持有序，中位数即其中间元素。

进一步，可考虑使用堆来优化时间复杂度。

建立一个小堆顶A，一个大堆顶B，各保存列表的一半元素。且规定：

  • A保存较大的一部分，长度为n/2，n为偶数；长度为n+1/2，n为奇数。

  • B保存较小的一部分，长度为n/2，n为偶数；长度为n-1/2，n为奇数。

中位数返回堆顶元素的和的一半即得。

增加元素

设元素总数为 N = m + n ，其中 m 和 n 分别为 A 和 B 中的元素个数。

  • 当 m=n，即N为偶数：需向A添加一个元素。实现方法：将新元素放入B，并将B的堆顶元素（大的数）放入A。
  • 当N为奇数，需向B添加一个元素（因为B少一个）。实现方法：将新元素放入A，并将A的堆顶元素（小的数）放入B。

代码实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
class MediumFinder{
    Queue<Integer> A,B;
    public MediumFinder(){
        A=new PriorityQueue<>();
        B=new PriorityQueue<>((x,y)->(y-x)); //java实现大堆顶。
    }
    public void addNum(int num){
        if(A.size()!=B.size()){
            A.add(num);
            B.add(A.poll());
        }else{
            B.add(num);
            A.add(B.poll());
        }
    }
    public double findMedium(){
        return A.size()!==B.size()? A.peek():(A.peek()+B.peek())/2.0;
    }
}