295. 数据流的中位数

发布于 2024-06-13  6 次阅读


中位数是有序整数列表中的中间值。如果列表的大小是偶数,则没有中间值,中位数是两个中间值的平均值。

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

实现 MedianFinder 类:

  • MedianFinder() 初始化 MedianFinder 对象。
  • void addNum(int num) 将数据流中的整数 num 添加到数据结构中。
  • double findMedian() 返回到目前为止所有元素的中位数。与实际答案相差 10-5 以内的答案将被接受。

示例 1:

输入
["MedianFinder", "addNum", "addNum", "findMedian", "addNum", "findMedian"]
[[], [1], [2], [], [3], []]
输出

[null, null, null, 1.5, null, 2.0]

解释

MedianFinder medianFinder = new MedianFinder();

medianFinder.addNum(1); // arr = [1]

medianFinder.addNum(2); // arr = [1, 2]

medianFinder.findMedian(); // 返回 1.5 ((1 + 2) / 2)

medianFinder.addNum(3); // arr[1, 2, 3]

medianFinder.findMedian(); // return 2.0

思路:使用两个堆,左边放比中位数小的数(大根堆),右边放比中位数大的数(小根堆),默认个数为奇数时左边数量比右边多一个,偶数则相同。在放入元素时,判断是大于还是小于堆顶元素,如果小于左边堆顶元素,放入左边,当左边数量比右边 + 1还大时,弹出堆顶元素放入右边。当数比右边堆顶元素大时,放入右边,当右边数量比左边大时,弹出堆顶元素放入左边。获取中位数时若左边数量大于右边,直接弹出堆顶元素即为中位数,反之取左右的堆顶元素计算平均值即为中位数。

class MedianFinder {

    PriorityQueue<Integer> min;
    PriorityQueue<Integer> max;

    public MedianFinder() {
        min = new PriorityQueue<>((a,b) -> (b - a));
        max = new PriorityQueue<>();
    }
    
    public void addNum(int num) {
        if(min.isEmpty() || num <= min.peek()){
            min.offer(num);
            if(max.size() + 1 < min.size()){
                max.offer(min.poll());
            }
        }else {
            max.offer(num);
            if(max.size() > min.size()){
                min.offer(max.poll());
            }
        }
    }
    
    public double findMedian() {
        if(min.size() > max.size()){
            return min.peek();
        }
        return (min.peek() + max.peek()) / 2.0;
    }
}