[LeetCode 295]数据流的中位数 Find Median from Data Stream

每次进行addNum后都sort一次会TLE

使用一个最大堆和最小堆
【STL】使用priority_queue构造堆
每次 addNum操作都改变堆一次
最大堆的size大于最小堆时,最大堆的堆顶元素就是中位数
最大堆的size小于最小堆时,最小堆的堆顶元素就是中位数
二者相等时,最大最小堆顶元素的平均值是中位数

C++代码:

#include 
#include 
#include 
using namespace std;

class MedianFinder
{
public:
    priority_queue, greater> small_heap;
    priority_queue, less> big_heap;

    /** initialize your data structure here. */
    MedianFinder()
    {
    }

    void addNum(int num)
    {
        if (big_heap.empty())
        {
            big_heap.push(num);
            return;
        }
        if (big_heap.size() == small_heap.size())
        {
            if (num < big_heap.top())
                big_heap.push(num);
            else
                small_heap.push(num);
        }

        else if (big_heap.size() > small_heap.size())
        {
            if (num <= big_heap.top())
            {
                small_heap.push(big_heap.top());
                big_heap.pop();
                big_heap.push(num);
            }
            else
                small_heap.push(num);
        }
        else
        {
            if (num >= small_heap.top())
            {
                big_heap.push(small_heap.top());
                small_heap.pop();
                small_heap.push(num);
            }
            else
                big_heap.push(num);
        }
    }

    double findMedian()
    {
        if (big_heap.size() == small_heap.size())
            return (double(big_heap.top()) + double(small_heap.top()) / 2);
        else if (big_heap.size() > small_heap.size())
            return double(big_heap.top());
        else
            return double(small_heap.top());
    }
};

int main()
{
    MedianFinder s;
    s.addNum(1);
    cout << s.findMedian() << endl;
    s.addNum(2);
    cout << s.findMedian() << endl;
    s.addNum(3);
    cout << s.findMedian() << endl;
    
    return 0;

}

解法来自79min处
https://www.bilibili.com/video/av29912609/?p=1

你可能感兴趣的:([LeetCode 295]数据流的中位数 Find Median from Data Stream)