数据流中的中位数

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

例如,

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

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

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

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

思路:

合适的数据结构是关键,使用一个有序数组,每次添加数据之后也是一个有序数组,这样,每次可以O(1)时间内读出中间数,添加数据的时间复杂度为O(n)。使用一个avl树,添加数据的时间复杂度为O(logN)。但是avl树实现起来比较复杂。考虑有序数组,中间数把数组分为两部分,左边部分的数都小于中间数,右边部分的数都大于中间数,至于左边部分的数是否有序不重要,右边也是。左边的数加上中间数,就是一个大根堆,右边的数加上中间数就是一个小根堆。用两个优先队列表示左右两边的数,那么添加数据的时间复杂度也是log(N),读取中间数的时间复杂度为O(1)。可以根据当前数据的下标来绝对放置数据,偶数时,把数据放在大根堆,奇数时把数据放在小根堆。当大根堆的根大于小根堆的根时,交换两个根。

代码:


class MedianFinder {
private: 
    priority_queue<int, vector<int>,less<int>> queMax;
    priority_queue<int, vector<int>,greater<int>> queMin;
    int index;
public:
    /** initialize your data structure here. */
    MedianFinder() {
        index = 0;
    }

    void addNum(int num) {
        if(index%2==0) {
            queMax.push(num);
        }
        else {
            queMin.push(num);
        }
        if(!queMax.empty() && !queMin.empty() && queMax.top() > queMin.top()) {
            int tmp = queMin.top();
            queMin.pop();
            queMin.push(queMax.top());
            queMax.pop();
            queMax.push(tmp);
        }
        index++;
    }

    double findMedian() {
        if(index%2==0) {
            if(!queMax.empty() && !queMin.empty()) {
                return (queMax.top()+queMin.top())/2.0;
            }
            return 0;
        }
        else {
            return queMax.top();
        }
    }
};