剑指 Offer（第2版）面试题 41：数据流的中位数

剑指 Offer（第2版）面试题 41：数据流的中位数

题目来源：LeetCode 295. 数据流的中位数

如何得到一个数据流中的中位数？

如果从数据流中读出奇数个数值，那么中位数就是所有数值排序之后位于中间的数值。

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

解法1：优先队列

我们用两个优先队列 minHeap 和 maxHeap 分别记录大于中位数的数和小于等于中位数的数。其中，maxHeap 是大顶堆，top 是其最大的元素；minHeap 是小顶堆，top 是其最小的元素。

void addNum(int num)：

当我们尝试添加一个数 num 到数据结构中，我们直接插入到 maxHeap 中。当 maxHeap.size() > minHeap.size() + 1 时，我们需要将 maxHeap.top() 移动到 minHeap 中。另外，当 maxHeap.top() > minHeap.top() 且 minHeap 不为空时，我们需要 swap(maxHeap.top(), minHeap.top()) 直至 maxHeap.top() <= minHeap.top()。

注意：

1. 在交换 2 个堆的元素的时候，一定要 先判断上面的小根堆中有没有元素。小根堆中没有元素是不能交换的。
2. 小根堆没有元素，就把大根堆中的 top 放到小根堆中。

double findMedian()：

• 当累计添加的数的数量为奇数时，maxHeap 中的数的数量比 minHeap 多一个，此时中位数为 maxHeap 的队头。
• 当累计添加的数的数量为偶数时，两个优先队列中的数的数量相同，此时中位数为它们的队头的平均值。

代码：

/*
 * @lc app=leetcode.cn id=295 lang=cpp
 *
 * [295] 数据流的中位数
 */

// @lc code=start
class MedianFinder
{
private:
    // 大顶堆，top 是其最大的元素，存储小于等于中位数的数
    priority_queue<int, vector<int>, less<int>> maxHeap;
    // 小顶堆，top 是其最小的元素，存储大于中位数的数
    priority_queue<int, vector<int>, greater<int>> minHeap;

public:
    MedianFinder() {}

    void addNum(int num)
    {
        maxHeap.push(num);
        if (maxHeap.size() > minHeap.size() + 1)
        {
            int x = maxHeap.top();
            minHeap.push(x);
            maxHeap.pop();
        }
        while (!minHeap.empty() && maxHeap.top() > minHeap.top())
        {
            int x = maxHeap.top(), y = minHeap.top();
            maxHeap.pop();
            maxHeap.push(y);
            minHeap.pop();
            minHeap.push(x);
        }
    }

    double findMedian()
    {
        double median;
        if ((maxHeap.size() + minHeap.size()) % 2 == 1)
            median = (double)maxHeap.top();
        else
            median = (maxHeap.top() + minHeap.top()) / 2.0;
        return median;
    }
};

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

复杂度分析：

时间复杂度：

• addNum：O(log⁡n)，其中 n 为累计添加的数的数量。
• findMedian：O(1)。

空间复杂度：O(n)，主要为优先队列的开销。

解法2：有序集合 + 双指针

我们也可以使用有序集合维护这些数。我们把有序集合看作自动排序的数组，使用双指针指向有序集合中的中位数元素即可。当累计添加的数的数量为奇数时，双指针指向同一个元素。当累计添加的数的数量为偶数时，双指针分别指向构成中位数的两个数。

当我们尝试添加一个数 num 到数据结构中，我们需要分情况讨论：

1. 初始有序集合为空时，我们直接让左右指针指向 num 所在的位置。
2. 有序集合为中元素为奇数时，left 和 right 同时指向中位数。如果 num 小于中位数，那么只要让 left 左移，否则让 right 右移。
3. 有序集合为中元素为偶数时，left 和 right 分别指向构成中位数的两个数。
   • 当 num 成为新的唯一的中位数，那么我们让 left 右移，right 左移，这样它们即可指向 num 所在的位置；
   • 当 num 大于等于 right 指向的值，那么我们让 left 右移即可；
   • 当 num 小于 right 指向的值，那么我们让 right 左移，注意到如果 num 恰等于 left 指向的值，那么 num 将被插入到 left 右侧，使得 left 和 right 间距增大，所以我们还需要额外让 left 指向移动后的 right。

代码：

// 有序集合 + 双指针

class MedianFinder
{
private:
    multiset<int> nums;
    multiset<int>::iterator left, right;

public:
    MedianFinder() : left(nums.end()), right(nums.end()) {}

    void addNum(int num)
    {
        const size_t n = nums.size();

        nums.insert(num);
        if (n == 0)
        {
            left = right = nums.begin();
        }
        else if (n % 2 == 1)
        {
            if (num < *left)
                left--;
            else
                right++;
        }
        else
        {
            if (num > *left && num < *right)
            {
                left++;
                right--;
            }
            else if (num >= *right)
                left++;
            else
            {
                right--;
                left = right;
            }
        }
    }

    double findMedian()
    {
        return (*left + *right) / 2.0;
    }
};

复杂度分析：

时间复杂度：

• addNum：O(log⁡n)，其中 n 为累计添加的数的数量。
• findMedian：O(1)。

空间复杂度：O(n)，主要为有序集合的开销。

一些进阶小 tips：