[LeetCode] Find Median from Data Stream

题目描述

Median is the middle value in an ordered integer list. If the size of the list is even, there is no middle value. So the median is the mean of the two middle value.

Examples:
[2,3,4] , the median is 3

[2,3], the median is (2 + 3) / 2 = 2.5

Design a data structure that supports the following two operations:

void addNum(int num) - Add a integer number from the data stream to the data structure.
double findMedian() - Return the median of all elements so far.
For example:

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

思路

用一个最大堆和一个最小堆维护所有数据,最小堆里的数据全部 >= 最大堆里的数据,
最小堆的size == 最大堆的size 或者 最小堆的size + 1 == 最大堆的size

这样的话,如果最大堆比最小堆多一个number,那么最大堆的堆顶一定是中位数,
如果最小堆和最大堆数目一样,那么median就是两个堆顶的平均数

判断每次加入number的大小,如果 > maxHeap.top(),那么加入到最小堆;
如果 <= maxHeap.top(),那么加入到最大堆;

然后对两个堆进行 balanced

C++代码

class MedianFinder {
    priority_queue<int, vector<int>, greater<int> > minheap;
    priority_queue<int, vector<int>, less<int> > maxheap;
public:

    // Adds a number into the data structure.
    void addNum(int num) {
        if (maxheap.empty()) {
            maxheap.push(num);
        }
        else {
            if (num <= maxheap.top()) {
                maxheap.push(num);
            }
            else {
                minheap.push(num);
            }
            // adjust two heaps
            if (minheap.size() + 1 < maxheap.size()) {
                int t = maxheap.top();
                maxheap.pop();
                minheap.push(t);
            }
            else if (minheap.size() > maxheap.size()) {
                int t = minheap.top();
                minheap.pop();
                maxheap.push(t);
            }
        }
    }

    // Returns the median of current data stream
    double findMedian() {
        if (maxheap.empty())
            return 0;
        else if (maxheap.size() > minheap.size())
            return maxheap.top();
        else
            return (maxheap.top() + minheap.top()) / 2.0;
    }
};

// Your MedianFinder object will be instantiated and called as such:
// MedianFinder mf;
// mf.addNum(1);
// mf.findMedian();