295. Find Median from Data Stream.md

295. Find Median from Data Stream

思路

有一个数据流，求中位数。

思路一、平衡二叉搜索树

维护一个始终有序的容器，而且插入的时间复杂度要尽可能低，所以可以考虑平衡二叉搜索树，在STL中就有现成的set和multiset，其内部实现为红黑树。这里我们使用multiset，因为数据可能重复，为了快速获取中位数，我们还要用一个始终指向中位数的指针mid（若元素数目为偶数我们规定mid指向较小的那个，或者用两个指针也是可以的），所以在插入新元素时需要合理调整mid：

• 若之前元素个数为奇数，插入新元素后变成偶数个：

  • 若num < *mid，则此时mid应该指向他前面的某个元素，所以应该将其前移一步，即mid = prev(mid)或者mid--；
  • 若num > *mid，此时mid刚好指向的就是两个中位数中较小的那个，所以mid保持不变；
  • 若num == *mid，由于STL规定插入到上界，所以和num > *mid情况是一样的，mid保持不变。

• 若之前元素个数为偶数，插入新元素后变成奇数个：

  • 若num < *mid，此时mid刚好就指向中位数，所以mid保持不变；
  • 若num >= *mid，此时mid指向的是中位数前面的那个元素，所以应该将其后移一步，即mid = next(mid)或mid++；

将上面的情况总结一下，就是

• 若之前的元素个数为奇数且num < *mid，则mid = prev(mid)或者mid--；
• 否则，若之前元素个数为偶数且num >= *mid，则mid = next(mid)或mid++；

最后求中位数就很方便了，见代码。

addNum的时间复杂度为O(logn)，findMedian的时间复杂度为O(1)。

注意：

• 迭代器一般（vector除外）没有加减操作，只能自加和自减，所以代码最后的*next(mid)不能写成*(mid+1)；
• 要想实现迭代器前移和后移可用 prev(it, n=1)和next(it, n=1)分别表示前移和后移n步。

思路二、大小堆

此题最经典的解法应该就是大小堆了，

• 用一个最大堆存放较小的那一半的元素，
• 用一个最小堆存放较大的那一半元素。

另外还需要保持两个堆元素均匀，这样就可以根据两个堆的堆顶元素在常数复杂度内求的中位数。

我们规定如果元素总数为奇数那么最大堆元素个数比最小堆元素个数多一个。那么插入元素算法为：

• 若之前元素个数为奇数，则最小堆需要新添加一个元素，这个元素是最大堆所有元素和新元素num中的最大值；
• 若之前元素个数为偶数，则最大堆需要新添加一个元素，这个元素是最小堆所有元素和新元素num中的最小值；

由于堆的插入和删除复杂度都是O(logn)，所以addNum时间复杂度为O(logn)；获取堆顶元素为常数复杂度，所以findMedian时间复杂度为O(1)。

C++

思路一

class MedianFinder {
private:
    int n = 0;
    multiset<int>st;
    multiset<int>::iterator mid;
public:
    /** initialize your data structure here. */
    MedianFinder() {}
    
    void addNum(int num) {
        st.insert(num);
        if(!n) mid = st.begin();
        else if(n & 1 && num < *mid) mid--; // 或 mid = prev(mid);
        else if(!(n & 1) && num >= *mid) mid++;// 或 mid = next(mid);
        n++;
    }
    
    double findMedian() {
        // *next(mid) 不能写成 *(mid + 1)
        return 0.5 * (*mid + (n & 1 ? *mid : *next(mid)));
    }
};

思路二

class MedianFinder {
private:
    int n = 0;
    priority_queue<int>max_heap;
    priority_queue<int, vector<int>, greater<int>>min_heap;
public:
    /** initialize your data structure here. */
    MedianFinder() {}
    
    void addNum(int num) {
        if(n & 1){
            if(num < max_heap.top()){
                max_heap.push(num);
                num = max_heap.top(); max_heap.pop();
            }
            min_heap.push(num);
        }
        else{
            if(!min_heap.empty() && num > min_heap.top()){                
                min_heap.push(num);
                num = min_heap.top(); min_heap.pop();
            }
            max_heap.push(num);
        }
        n++;
    }
    
    double findMedian() {
        return 0.5 * (max_heap.top() + (n & 1 ? max_heap.top() : min_heap.top()));
    }
};