Leetcode刷题笔记——剑指 Offer 41. 数据流中的中位数（困难）

---

题目描述

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

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

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

解法

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

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

借助 堆 可进一步优化时间复杂度。

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

• $A$ 保存 较大 的一半，长度为 $\frac{N}{2}$（ $N$ 为偶数）或 $\frac{N+1}{2}$（ $N$ 为奇数）；
• $B$ 保存 较小 的一半，长度为 $\frac{N}{2}$（ $N$ 为偶数）或 $\frac{N-1}{2}$（ $N$ 为奇数）；

随后，中位数可仅根据 $A,B$ 的堆顶元素计算得到。

算法流程：

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

addNum(num) 函数：

1. 当 $m=n$ (即 $N$ 为 偶数）：需向 $A$ 添加一个元素。实现方法：将新元素 $num$ 插入至 $B$ ，再将 $B$ 堆顶元素插入至 $A$ ；
2. 当 $m\ne n$（即 $N$ 为 奇数）：需向 $B$ 添加一个元素。实现方法：将新元素 $num$ 插入至 $A$ ，再将 $A$ 堆顶元素插入至 $B$ ；

假设插入数字 $num$ 遇到情况 1. 。由于 $num$ 可能属于 “较小的一半” （即属于 $B$ ），因此不能将 $nums$ 直接插入至 $A$ 。而应先将 $num$ 插入至 $B$ ，再将 $B$ 堆顶元素插入至 $A$ 。这样就可以始终保持 $A$ 保存较大一半、 $B$ 保存较小一半。

findMedian() 函数：

1. 当 $m=n$（ $N$ 为 偶数）：则中位数为 (( $A$ 的堆顶元素 + $B$ 的堆顶元素 )$/2$)。
2. 当 $m\ne n$ （ $N$ 为 奇数）：则中位数为 $A$ 的堆顶元素。

复杂度分析

• 时间复杂度：
  • 查找中位数 $O(1)$ ： 获取堆顶元素使用 $O(1)$ 时间；
  • 添加数字 $O(\log N)$ ： 堆的插入和弹出操作使用 $O(\log N)$ 时间。
• 空间复杂度 $O(N)$ ： 其中 $N$ 为数据流中的元素数量，小顶堆 $A$ 和大顶堆 $B$ 最多同时保存 $N$ 个元素。

C++代码实现

class MedianFinder {
public:
    // 最大堆，存储左边一半的数据，堆顶为最大值
    priority_queue<int, vector<int>, less<int>> maxHeap;
    // 最小堆， 存储右边一半的数据，堆顶为最小值
    priority_queue<int, vector<int>, greater<int>> minHeap;
    /** initialize your data structure here. */
    MedianFinder() {
    }

    // 维持堆数据平衡，并保证左边堆的最大值小于或等于右边堆的最小值
    void addNum(int num) {
        /*
         * 当两堆的数据个数相等时候，左边堆添加元素。
         * 采用的方法不是直接将数据插入左边堆，而是将数据先插入右边堆，算法调整后
         * 将堆顶的数据插入到左边堆，这样保证左边堆插入的元素始终是右边堆的最小值。
         * 同理左边数据多，往右边堆添加数据的时候，先将数据放入左边堆，选出最大值放到右边堆中。
         */
        if (maxHeap.size() == minHeap.size()) {
            minHeap.push(num);
            int top = minHeap.top();
            minHeap.pop();
            maxHeap.push(top);
        } else {
            maxHeap.push(num);
            int top = maxHeap.top();
            maxHeap.pop();
            minHeap.push(top);
        }
    }
    
    double findMedian() {
        if (maxHeap.size() == minHeap.size()) {
            return (maxHeap.top()+minHeap.top())*1.0/2;
        } else {
            return maxHeap.top()*1.0;
        }
    }
};

参考链接

[1] https://leetcode-cn.com/problems/shu-ju-liu-zhong-de-zhong-wei-shu-lcof/solution/mian-shi-ti-41-shu-ju-liu-zhong-de-zhong-wei-shu-y/