leetcode题295.数据流的中位数

题目

295. 数据流的中位数

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

  • 例如 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
2
3
4
5
6
7
8
9
10
11
12
13
14
示例 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

提示:

  • -10^5 <= num <= 10^5;
  • 在调用 findMedian 之前，数据结构中至少有一个元素;
  • 最多 5 * 104 次调用 addNum 和 findMedian;

思路

来自k神的题解 https://leetcode.cn/problems/find-median-from-data-stream/solutions/2361972/295-shu-ju-liu-de-zhong-wei-shu-dui-qing-gmdo/ ，用两个堆来解决这个问题。

首要的思想是，我们需要将数组流存入一个有序的数组，再确定中位数。不用堆的情况下，我们可以利用插入排序的思想，将新数据插入到数组中的正确位置。这涉及到O(N)的数组中元素的移动，以及利用二分查找数组的插入位置O(LogN)。总的时间复杂度是O(logN+N)

使用两个堆，就可以将整体的时间复杂度控制在O(logN)上，一定程度上能提高效率。

  • 小堆A用来存放数组大的那一部分数据（你可以理解为堆内元素是倒序，堆顶是最小的数）
  • 大堆B用来存放数组小的那一部分的数据（堆内元素是正序，堆顶是最大的数）
  • A中的所有数据都必须大于B

此时小堆A对着大堆B，就构成了一个完整的“有序数组”。

1
2
小堆A [6,5,4]   [3,2,1] 大堆B
   小堆堆顶4    大堆堆顶3

插入的时候，约定小堆A的元素数量应该大于大堆B，分为奇数和偶数的情况

  • A和B的大小相等，说明当前数据流数量是偶数，往A中插入元素
    • 需要先将待插入元素num和B的堆顶元素比较（确保A中是数组大的那一部分的数据）
    • 如果num大于B的堆顶元素，则直接插入A；
    • 如果num小于B的堆顶元素，则插入B后，往A内插入B的堆顶元素，再弹出B的堆顶元素；
  • A的元素个数大于B，说明当前数据流数量是奇数，往B中插入元素
    • 将插入元素和A的堆顶元素比较
    • 如果num大于A的堆顶元素，则将A的堆顶元素插入B，弹出A的堆顶元素，再将num插入A；
    • 如果num小于A的堆顶元素，则将num直接插入B

查询中位数的时候，也分为奇数和偶数的情况，注意题目要求的返回值是double

  • A和B的元素个数相等，说明是偶数，中位数为(A堆顶元素+B堆顶元素)/2.0；
  • A和B的元素个数不相等（肯定是A的元素多），说明是奇数，中位数为A的堆顶元素*1.0。

思路会了，写代码就不难了。

代码

注意cpp中优先级队列默认是大堆，且std库中提供了less和greater作为堆的比较函数

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
class MedianFinder {
    // 注意优先级队列C++中默认的比较函数是less，存放的是大堆
    priority_queue<int> _queBig; // 大堆，存放小半边数据
    priority_queue<int, vector<int>, std::greater<int>>
        _queSm; // 小堆，存放大半边数据
public:
    MedianFinder() {}

    // 保证较大的一半数据更多
    // 奇数的时候较大的一半堆顶就是中位数
    // 偶数的时候，较小+较大堆顶/2就是中位数
    void addNum(int num) {
        if (_queBig.size() == _queSm.size()) {
            // 注意sm存放的是大的那一半
            // _queSm.push(num);
            // error！不能直接插入!
            // 先把num插入到big里面，再把big里面最大的那个数插入到sm里面！
            // 因为需要保证有序，如果直接插入，此时num的值可能比big里面的堆顶元素小！
            _queBig.push(num);
            _queSm.push(_queBig.top());
            _queBig.pop();
            return;
        }
        // 大的那一半数据更多
        // 注意sm存放的是大的那一半
        if (_queBig.size() < _queSm.size()) {
            // 如果当前数比大堆堆顶的更大，则需要放到大的那一半中
            if (num > _queSm.top()) {
                _queBig.push(_queSm.top());
                _queSm.pop();
                _queSm.push(num);
                return;
            }
            // 其他情况都往小的那一半里面插
            _queBig.push(num);
        }
        return;
    }

    double findMedian() {
        if (_queBig.size() == _queSm.size()) {
            return double((_queBig.top() + _queSm.top()) / 2.0);
        }
        // sm才是大的那一半
        return double(_queSm.top() * 1.0);
    }
};

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

image-20240330163821849