PROBLEM

剑指 Offer 41. 数据流中的中位数

难度 困难

MY ANSWER

一个大顶堆一个小顶堆，分别存放比中位数小和大的数。addnum时间复杂度O(logn)，空间复杂度O(n)，findmedian时间复杂度O(1)，空间复杂度O(1)。

class MedianFinder {
public:
    priority_queue<int> small;
    priority_queue<int, vector<int>, greater<int>> big;
    double median = INT_MIN;
    MedianFinder() {
    }
    
    void addNum(int num) {
        if(median == INT_MIN) {
            median = num;
        }
        if(num <= median) {
            if(small.size() <= big.size()) {
                small.push(num);
            }
            else {
                big.push(small.top());
                small.pop();
                small.push(num);
            }
        }
        else {
            if(small.size() >= big.size()) {
                big.push(num);
            }
            else {
                small.push(big.top());
                big.pop();
                big.push(num);
            }
        }
        if(small.size() == big.size()) {
            median = (double)(small.top() + big.top()) / 2;
        }
        else if(small.size() < big.size()) {
            median = big.top();
        }
        else {
            median = small.top();
        }
    }
    
    double findMedian() {
        return median;
    }
};

/**
 * Your MedianFinder object will be instantiated and called as such:
 * MedianFinder* obj = new MedianFinder();
 * obj->addNum(num);
 * double param_2 = obj->findMedian();
 */

BETTER SOLUTION

同样一个大顶堆一个小顶堆，不需要每次addnum时都计算中位数，然后再将数与中位数比较。每插入一个数，判断两个堆大小是否一样。若一样，插入大顶堆，将大顶堆堆顶插入小顶堆，保证小顶堆所有数大于等于大顶堆。若不一样，插入小顶堆，将小顶堆堆顶插入大顶堆，保证大顶堆所有数小于等于小顶堆。这样还能保证两个堆的大小只相差1。计算中位数时，若两堆大小不等，直接返回小顶堆堆顶；若相等，计算两堆顶之和的一半。时间空间复杂度同上。

class MedianFinder {
public:
    /** initialize your data structure here. */
    priority_queue<int, vector<int>, less<int> > maxheap;
    priority_queue<int, vector<int>, greater<int> > minheap;

    MedianFinder() {

    }
    
    void addNum(int num) {
        if(maxheap.size() == minheap.size()) {
            maxheap.push(num);
            minheap.push(maxheap.top());
            maxheap.pop();
        }
        else {
            minheap.push(num);
            maxheap.push(minheap.top());
            minheap.pop();
        }
    }
    
    double findMedian() {
        int maxSize = maxheap.size(), minSize = minheap.size();
        int mid1 = maxheap.top(), mid2 = minheap.top();
        return maxSize == minSize ? ((mid1 + mid2) * 0.5) : mid2;
    }
};

SUMMARY

  • 小顶堆：priority_queue<int, vector<int>, greater<int> > minheap，默认为大顶堆。
  • 频繁的插入、删除考虑使用优先队列。