Find Median from Data Stream 295

题目链接：https://leetcode.com/problems/find-median-from-data-stream/

 

题目描述： 设计一个类，包含两种操作

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

操作 addNum向数列中添加数字

操作findMedian返回数列中的中位数，当数列个数是偶数的时候返回中间两个数字的平均数

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

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

 

题目分析：

假如当前数列中的中位数是n，那么n将数列分为两部分A，B。A中数字都小于n，B中数字都大于n。

假如再添加一个数字m，

如果m>n那么中位数应该是B中最小的数字，或者B中最小的数字和n的平均数

如果m<n那么中位数应该是A中最大的数字，或者A中最大的数字和n的平均数。

从A中取出最大数字和重B中取出最小数字可以用大顶推和小顶堆实现。

 

设计一个大顶堆和一个小顶堆，大顶堆中的元素都小于小顶堆中的元素，两个堆中元素个数或者相等或者大顶堆中元素比小顶堆重元素多一个。

这样当两个堆中元素个数不相同的时候大顶堆中的最大元素就是中位数，当两个对中元素个数相同的时候，大顶堆中最大的元素和小顶堆中最小的元素的平均数，就是我们要求的结果。

 

 

算法实现：

#include <vector>
#include <map>
#include <cstring>
#include <cstdlib>
#include <cstdio>
#include <algorithm>
#include <iostream>
#include <stack>

using namespace std;

class MedianFinder {
public:

    MedianFinder() {
        small_heap.push_back(0);
        big_heap.push_back(0);
    }
    // Adds a number into the data structure.
    void addNum(int num) {
        if (!big_heap[0]) {
            ++big_heap[0];
            big_heap.push_back(num);
            return;
        }

        if (num <= big_heap[1]) {
            if (big_heap[0] <= small_heap[0]) {
                ++big_heap[0];
                big_heap.push_back(num);
                push_up(big_heap, big_heap[0], big);
            }
            else {
                ++small_heap[0];
                small_heap.push_back(big_heap[1]);
                push_up(small_heap, small_heap[0], small);
                big_heap[1] = num;
                push_down(big_heap, 1, big);
            }
        }
        else {
            if (small_heap[0] < big_heap[0]) {
                ++small_heap[0];
                small_heap.push_back(num);
                push_up(small_heap, small_heap[0], small);
            }
            else {
                if (num > small_heap[1]) swap(small_heap[1], num);
                push_down(small_heap, 1, small);
                ++big_heap[0];
                big_heap.push_back(num);
                push_up(big_heap, big_heap[0], big);
            }
        }
    }

    // Returns the median of current data stream
    double findMedian() {
        //printf("find_median: big_heap num %d, small_heap num %d\n", big_heap[0], small_heap[0]);
        //printf("big_heap: "); print_vector(big_heap);
        //printf("small_heap: "); print_vector(small_heap);
        if (big_heap[0] > small_heap[0]) return big_heap[1];
        return (double)(big_heap[1] + small_heap[1]) / 2;
    }

private:
    enum heap_type {small = -1, big = 1};
    vector<int> small_heap;
    vector<int> big_heap;

    void push_down(vector<int> &heap, int cur, heap_type type) {
        int l = cur * 2;
        int r = l + 1;
        int p = cur;
        if (l <= heap[0] && heap[p]*type < heap[l]*type) p = l;
        if (r <= heap[0] && heap[p]*type < heap[r]*type) p = r;

        if (p == cur) return;
        swap(heap[cur], heap[p]);
        
        push_down(heap, p, type);
    }
    
    void push_up(vector<int> &heap, int cur, heap_type type) {
        if (cur == 1) return ;
        int fa = cur >> 1;
        if (heap[fa]*type >= heap[cur]*type)
            return ;
        swap(heap[fa], heap[cur]);
        push_up(heap, fa, type);
    }
    void print_vector(vector<int> vet) {
        for (int i = 0; i < vet.size(); i++) {
            cout<<vet[i]<<"\t";
        }
        cout<<endl;
    }
};

int main()
{
    MedianFinder mf;
    mf.addNum(1);
    mf.addNum(2);
    mf.addNum(2);
    mf.addNum(3);
    mf.addNum(4);
    mf.addNum(5);
    double ret = mf.findMedian();
    cout<<ret<<endl;

    return 0;
}