295.Find Median from Data Stream

Tags: Hard Heap Design

Links:

---

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.

For example,

[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.

Example:

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

Follow up:

1. If all integer numbers from the stream are between 0 and 100, how would you optimize it?
2. If 99% of all integer numbers from the stream are between 0 and 100, how would you optimize it?

---

class MedianFinder {
    priority_queue<int, vector<int>, greater<int>> right;
    priority_queue<int> left;
    int n;
    double mid;

public:
    /** initialize your data structure here. */
    MedianFinder() {
        std::ios_base::sync_with_stdio(false);
        cin.tie(NULL);
        cout.tie(NULL);

        n = 0;
        mid = 0.0;
    }

    void addNum(int num) {
        if (n == 0) {
            ++n;
            mid = num;
            return;
        }

        if (n & 1) { //原先的数据个数为奇数
            ++n;
            if (num * 1.0 >= mid) {
                left.push(mid);
                right.push(num);
            }
            else {
                left.push(num);
                right.push(mid);
            }
            mid = (left.top() + right.top()) * 1.0 / 2;
        }
        else { //原有偶数个数据
            ++n;
            if (num * 1.0 >= mid) {
                right.push(num);
                mid = right.top();
                right.pop();
            }
            else {
                left.push(num);
                mid = left.top();
                left.pop();
            }
        }
    }

    double findMedian() {
        return mid;
    }
};

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

两个堆，一个大根堆，一个小根堆。对于数据大根堆维护左半部分，小根堆维护右半部分，另外用mid代表中位数。

插入数据的时候考虑两种情况：

• 没有数据的时候，那么中位数就是插入的数据本身，计数器+1
• 有数据
  • 原来的数据个数是奇数：插入一个数据，数据总数就是偶数了，那么去比较插入的数据和mid的值，注意要先把插入的数据转为浮点数再比较，然后两个数据分别插入左右两个堆，更新mid。
  • 原来的数据个数是偶数：插入一个数据，数据总数变为奇数，并且中位数就是所有数据中的某一个数，所以用mid来存储这个数据，这样左右两个堆的数据就相同了，那么考虑插入的数据大于mid，则把新插入的数据放入右堆，然后让mid为右堆的最小值；如果插入的数据小于mid，那么数据插入左堆，然后让mid为左堆的最大值。

扩展部分：

• 如果所有数据都是再0-100之间，如何维护？很显然，用线段树维护。
• 如果99%的数据在0-100之间，如何维护？显然大于100的数据不会很多，可以把多出100的数据用mutiset维护，在0-100之间的还是用线段树维护，查找的时候基本也是一致的。