LeetCode 295. 数据流的中位数

Posted 2021-01-10 hlk09

中位数是有序列表中间的数。如果列表长度是偶数，中位数则是中间两个数的平均值。

例如，

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

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

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

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

示例：

addNum(1)
addNum(2)
findMedian() -> 1.5
addNum(3) 
findMedian() -> 2

这个题太过巧妙，我是直接学习LeetCode英文版的解答的。在这个题中，使用了两个堆（优先队列）。其中，一个堆为大顶堆，一个堆为小型堆。大顶堆存数组左边的数据，小顶堆存数组右边的数据。那么，小顶堆堆顶和大顶堆对顶是唯一提供中位数的可能来源。

为了保持大顶堆和小顶堆的平衡，设定小顶堆比大顶堆最多存储多一个元素。

至于添加操作，首先是把添加元素放入大顶堆。再把大顶堆的最大元素转移到小顶堆。如果小顶堆比大顶堆大，再把小顶堆元素转移到大顶堆。

class MedianFinder {
    priority_queue<int> lo;                              // max heap
    priority_queue<int, vector<int>, greater<int>> hi;   // min heap

public:
    // Adds a number into the data structure.
    void addNum(int num)
    {
        lo.push(num);                                    // Add to max heap

        hi.push(lo.top());                               // balancing step
        lo.pop();

        if (lo.size() < hi.size()) {                     // maintain size property
            lo.push(hi.top());
            hi.pop();
        }
    }

    // Returns the median of current data stream
    double findMedian()
    {
        return lo.size() > hi.size() ? (double) lo.top() : (lo.top() + hi.top()) * 0.5;
    }
};