剑指 Offer 41. 数据流中的中位数

题目

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

例如，

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

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

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

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

示例 1：

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

示例 2：

输入： [“MedianFinder”,“addNum”,“findMedian”,“addNum”,“findMedian”]
[[],[2],[],[3],[]]
输出：[null,null,2.00000,null,2.50000]

限制：

• 最多会对 addNum、findMedian 进行 50000 次调用。

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/shu-ju-liu-zhong-de-zhong-wei-shu-lcof

解题思路

首先理解题意,数据流中的中位数前提都是对于有序数组

• 数组大小为奇数,就是中间的那一个数
• 数组大小为偶数,就是中间的两个数的平均值

解题思路就是使用两个堆来维护一个有序数组

• 一个大顶堆加上一个小顶堆可以当做一个有序数组

• 大顶堆作为数组的前半段,存储数组中较小的元素

• 小顶堆作为数组的后半段,存储数组中较大的元素

• 当数组的大小为偶数时,大顶堆和小顶堆的大小相等

• 当数组的大小为奇数时,可以约定奇数放在大顶堆还是小顶堆,这个都随意,代码不同

• 创建一个大顶堆a和小顶堆b

• 假设大顶堆a的大小为n,小顶堆b的大小为m,我们约定

  • 数组大小为偶数时n == m,
  • 数组大小为奇数时n == m + 1

• 添加元素时

  • 当n != m时,此时说明数组大小为奇数,添加一个元素变为偶数,则需要转换成n == m,即小顶堆元素+1
    • 小顶堆存储数组的后半段,存储数组中较大的元素,待添加元素和大顶堆堆顶元素中较大的元素入大顶堆
    • 待添加元素,入堆大顶堆
    • 大顶堆堆顶元素入堆小顶堆
  • 当n == m时,此时说明数组大小为偶数,添加一个元素变为奇数,则需要转换成n == m + 1,即大顶堆元素+1
    • 大顶堆存储数组的前半段,存储数组中较小的元素,待添加元素和小顶堆堆顶元素中较小的元素入小顶堆
    • 待添加元素,入堆小顶堆
    • 小顶堆堆顶元素入堆大顶堆

• 获取中位数时

  • 当n != m时,直接返回大顶堆的堆顶元素
  • 当n == m时,返回大顶堆堆顶元素和小顶堆堆顶元素,取平均值

代码

class MedianFinder {

    // 维护两个堆
    PriorityQueue<Integer> a;
    PriorityQueue<Integer> b;
    /** initialize your data structure here. */
    public MedianFinder() {
        // 其中a为大顶堆,b为小顶堆
        // 大顶堆存储有序数组中较小的元素,小顶堆存储有序数组中较大的元素
        // 可以理解为大顶堆a是有序数组的前半段,小顶堆b是有序数组的后半段
        a = new PriorityQueue<>((v1, v2) -> v2 - v1);
        b = new PriorityQueue<>();
    }
    
    public void addNum(int num) {
        // 假设a堆大小为n,b堆大小为m,约定n == m,或者n == m + 1
        if (a.size() != b.size()) {
            // 当a堆和b堆大小不相等时,依据约定此时n == m + 1
            // 转移成n == m,则此时需要n不变,m + 1
            // b堆中保存数组中较大的元素,所以需要添加的是num和大顶堆a堆顶元素中的最大值
            // 元素加入a堆,然后把a堆堆顶的元素,即最大值入b堆
            a.add(num);
            b.add(a.poll());
        } else {
            // 当a堆和b堆大小相等时,依据约定此时n == m
            // 转移成n == m + 1,则此时需要m不变,n + 1
            // a堆中保存数组中较小的元素,所以需要添加的是num和小顶堆b堆顶元素中的最小值
            // 元素加入b堆,然后把b堆堆顶的元素,即最小值入a堆
            b.add(num);
            a.add(b.poll());
        }
    }
    
    public double findMedian() {
        if (a.size() != b.size()) {
            // a堆和b堆大小不相等,直接取a堆堆顶元素返回
            return a.peek();
        } else {
            // a堆和b堆大小相等,取两个堆顶元素相加再除以2.0
            return (a.peek() + b.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();
 */

复杂度

• 时间复杂度:
  • 添加: O(logn),堆的插入和弹出操作使用 O(\log N)O(logN) 时间
  • 查找:O(1),获取堆顶元素使用O(1)时间
• 空间复杂度:O(n),维护数组元素