【大小堆】【leetcode】【困难】295. 数据流的中位数

题目：

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

例如，

[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

原题地址：

295. 数据流的中位数

解题思路：大小堆

定义2个堆，大顶堆small，小顶堆big，并且保证small.size <= big.size，且max(small) <= min(big)。

定义变量totoal，记录插入的数量。当插入前为奇数时，small.size<big.size，为偶数时small.size==big.size。

新插入的数据如果比big.top()小，则插入small中，如果此时small中数量多于big，则将small堆顶数据挪至big中；如果插入big，如果插入前big数量比small多，则将big堆顶数据挪至small中。

c++堆由priority_queue实现，代码如下：

class MedianFinder {
public:
    /** initialize your data structure here. */
    MedianFinder() {
        total = 0;
    }
    
    void addNum(int num) {
        if (total == 0) {
            big.push(num);
            ++total;
            return;
        }
        
        int mid = big.top();
        if (num < mid) {
            small.push(num);// 插入small
            if (total % 2 == 0) {
                big.push(small.top());
                small.pop();
            }
        } else {
            big.push(num);// 插入big
            if (total % 2 == 1) {
                small.push(big.top());
                big.pop();
            }
        }
        ++total;
    }
    
    double findMedian() {
        if (total % 2 == 1) return big.top();
        else return ((double)small.top() + (double)big.top()) / 2;
    }
private:
    int total;
    // 保证：small.size() <= big.size()
    priority_queue<int> small; // 大顶堆
    priority_queue<int, vector<int>, greater<int> > big; // 小顶堆
};

也可用multiset实现，代码如下：

class MedianFinder {
public:
    /** initialize your data structure here. */
    MedianFinder() {
        total = 0;
    }
    
    void addNum(int num) {
        if (total == 0) {
            big.insert(num);
            ++total;
            return;
        }
        
        multiset<int>::iterator mid = big.begin();
        if (num < *mid) {
            small.insert(num);// 插入small
            if (total % 2 == 0) {
                multiset<int>::reverse_iterator it = small.rbegin();
                big.insert(*it);
                small.erase(small.find(*it));
            }
        } else {
            big.insert(num);// 插入big
            if (total % 2 == 1) {
                small.insert(*mid);
                big.erase(mid);
            }
        }
        ++total;
    }
    
    double findMedian() {
        if (total % 2 == 1) return *big.begin();
        else return ((double)*small.rbegin() + (double)*big.begin()) / 2;
    }
private:
    int total;
    // 保证：small.size() <= big.size()
    multiset<int> small;
    multiset<int> big;
};