[劍指Offer] 數據流中的中位數

題目描述

如何得到一個數據流中的中位數？如果從數據流中讀出奇數個數值，那么中位數就是所有數值排序之后位於中間的數值。如果從數據流中讀出偶數個數值，那么中位數就是所有數值排序之后中間兩個數的平均值。

 

對於數據流，對應的就是在線算法了，一道很經典的題目就是在1億個數中找到最大的前100個數，這是一道堆應用題，找最大的前100個數，那么我們就創建一個大小為100的最小化堆，每來一個元素就與堆頂元素比較，因為堆頂元素是目前前100大數中的最小數，前來的元素如果比該元素大，那么就把原來的堆頂替換掉。

那么對於這一道題呢？如果單純的把所有元素放到一個數組里，每次查找中位數最快也要O(n)，綜合下來是O(n^2)的復雜度。我們可以利用上面例子中的想法，用一個最大堆來維護當前前n/2小的元素，那么每次找中位數只到取出堆頂就可以了。但是，有一個問題，數據要動態增長，有可能之前被替換掉的元素隨着元素的增加又跑回來了，所以我們不能單純得向上題一樣把元素丟掉，我們可以再用一個最小化堆來存前n/2大的元素。

#include <bits/stdc++.h>
using namespace std;

class Solution {
    private:
        vector<int> min; //數組中的后一半元素組成一個最小化堆
        vector<int> max; //數組中的前一半元素組成一個最大化堆
    public:
        void Insert(int num) {
            if(((min.size()+max.size()) & 1) == 0) {  //偶數數據的情況下，則在最小堆中插入元素
                if(max.size() > 0 && num < max[0]) {
                    max.push_back(num);
                    push_heap(max.begin(), max.end(), less<int>());
                    num=max[0];
                    pop_heap(max.begin(), max.end(), less<int>());
                    max.pop_back();
                }
                min.push_back(num); //把前一半找到的最大值放到后一半中
                push_heap(min.begin(), min.end(), greater<int>());
            } else {
                if(min.size() > 0 && num > min[0]) {   //奇數數據的情況下，則在最大堆中插入元素
                    min.push_back(num);
                    push_heap(min.begin(), min.end(), greater<int>());
                    num=min[0];
                    pop_heap(min.begin(), min.end(), greater<int>());
                    min.pop_back(); 
                }
                max.push_back(num); //把后一半找到的最大值放到前一半中
                push_heap(max.begin(), max.end(), less<int>());
            }
        }

        double GetMedian() { 
            int size=min.size() + max.size();
            if(size==0) return -1;
            double median = 0;
            if((size&1) != 0) {
                median = (double) min[0];
            } else {
                median = (double) (max[0] + min[0]) / 2;
            }
            return median;
        }
};

int main() {
    Solution s;
    vector<int> v{5,2,3,4,1,6,7,0,8};
    for (int i = 0; i < v.size(); ++i) {
        s.Insert(v[i]);
        cout << s.GetMedian() << endl;
    }
    return 0;
}

也可以使用multiset來簡化編程，lintcode上也有原題。

http://www.lintcode.com/en/problem/data-stream-median/

class Solution {
public:
    /**
     * @param nums: A list of integers.
     * @return: The median of numbers
     */
    vector<int> medianII(vector<int> &nums) {
        // write your code here
        multiset<int> left, right;
        vector<int> res;
        bool flag = true;
        for (int n : nums) {
            int tmp = n;
            if (flag) {
                if (!right.empty() && n > *right.begin()) {
                    right.insert(n);
                    tmp = *right.begin();
                    right.erase(right.find(tmp));
                }
                left.insert(tmp);
            } else {
                if (!left.empty() && n < *left.rbegin()) {
                    left.insert(n);
                    tmp = *left.rbegin();
                    left.erase(left.find(tmp));
                }
                right.insert(tmp);
            }
            flag = !flag;
            res.push_back(*left.rbegin());
        }
        return res;
    }
};

還有一道是求滑動窗口中的中位數，其實是基於同樣的思想。只是在窗口滑動時，會有元素滑出窗口，所以在插入新的元素之前先要把滑出窗口的元素刪除掉。

http://www.lintcode.com/en/problem/sliding-window-median/

class Solution {
public:
    /**
     * @param nums: A list of integers.
     * @return: The median of the element inside the window at each moving
     */
    vector<int> medianSlidingWindow(vector<int> &nums, int k) {
        // write your code here
        vector<int> res;
        if (k > nums.size() || k == 0) return res;
        multiset<int> left, right;
        //init heaps by first kth elements in nums
        for (int i = 0; i < k; ++i) {
            left.insert(nums[i]);
        }
        while (left.size() > (k + 1) / 2) {
            right.insert(*left.rbegin());
            left.erase(left.find(*left.rbegin()));
        }
        res.push_back(*left.rbegin());
        //slide window
        for (int i = k; i < nums.size(); ++i) {
            //delete the leftmost element in window from heaps
            if (nums[i-k] > res.back()) right.erase(right.find(nums[i-k]));
            else left.erase(left.find(nums[i-k]));
            //insert new element into heaps
            if (!left.empty() && nums[i] <= *left.rbegin()) left.insert(nums[i]);
            else right.insert(nums[i]);
            //adjust heaps so that the left heap contains (k + 1) / 2 elements
            while (left.size() < (k + 1) / 2) {
                left.insert(*right.begin());
                right.erase(right.begin());
            }
            while (left.size() > (k + 1) / 2) {
                right.insert(*left.rbegin());
                left.erase(left.find(*left.rbegin()));
            }
            res.push_back(*left.rbegin());
        }
        return res;
    }
};