[LeetCode] Find Median from Data Stream | 动态计算数据流的中位数

https://leetcode.com/problems/find-median-from-data-stream/#/description

维护两个堆,一个最小堆,用于存比中位数大的数;一个最大堆,用于存比中位数小的数。关键在于使得两个堆的大小之差不超过1。

class MedianFinder {
public:
    /** initialize your data structure here. */
    MedianFinder() {

    }
    
    void addNum(int num) {
        if (median.empty()) {
            median.push_back(num);
            return;
        }
        if ((double)num <= median[0]) {
            max_heap.push(num);
        } else {
            min_heap.push(num);
        }
        if (max_heap.size() > min_heap.size() + 1) {
            min_heap.push(median[0]);
            median[0] = max_heap.top();
            max_heap.pop();
        } else if (min_heap.size() > max_heap.size() + 1) {
            max_heap.push(median[0]);
            median[0] = min_heap.top();
            min_heap.pop();
        }
    }
    
    double findMedian() {
        if (max_heap.size() == min_heap.size()) {
            return (double) median[0];
        } else if (max_heap.size() < min_heap.size()) {
            return ((double) median[0] + min_heap.top()) / 2;
        } else {
            return ((double) median[0] + max_heap.top()) / 2;
        }
    }

private:
    priority_queue<int> max_heap;
    priority_queue<int, vector<int>, greater<int>> min_heap;
    vector<int> median;
};

/**
 * Your MedianFinder object will be instantiated and called as such:
 * MedianFinder obj = new MedianFinder();
 * obj.addNum(num);
 * double param_2 = obj.findMedian();
 */
posted @ 2017-06-19 20:39  mioopoi  阅读(186)  评论(0编辑  收藏  举报