NC131 数据流中的中位数

题目描述:

设计一种数据结构,支持快速获取中位数。

1. 暴力解法

直接用一个数组承载,每次获取中位数的时候排序,再取中间点。

注意这里要求的是浮点数,可以用整数*1.0进行类型转换。

class Solution {
public:
    vector<int> array;
    void Insert(int num) {
        array.push_back(num);
    }

    double GetMedian() { 
        sort(array.begin(), array.end());
        int n = array.size();
        
        if (n&1) return array[n/2] * 1.0;
        else return (array[n/2-1] + array[n/2]) * 1.0 / 2.0;
    }

};
  • 时间复杂度: Insert的时间复杂度是O(1)O(1),GetMedian用到了一轮排序,是 O(nlogn)O(nlogn)
  • 空间复杂度:O(n)O(n)

2. 使用两个堆调整(正解)

其实问题要求的是中位数,我们对整个数组排序显然是冗余的了。我们需要把中间的部分“漏出来”,才能动态维护这个中位数,因此我们想到的数据结构是

我们可以想到用堆维护数据流上半部分最大值,下半部分最小值,并且使得小根堆的最大值小于大根堆的最小值, 并且两个堆是平衡的,这样才能保证每次两个堆的堆顶可以算出中位数。实现时需要注意,新添加进去的数据需要判断该和哪个堆比较,并维护当前要加入的元素是第奇数还是偶数个。

class Solution {
public:
    priority_queue<int> max_heap;
    // 小根堆,保证比大根堆的个数多1个或等于
    priority_queue<int, vector<int>, greater<int>> min_heap;
    
    // 初始化时,当前是第0个元素,0是偶数。
    bool isEven = true;
    
    void Insert(int num) {
        if (isEven) {
            // 如果是偶数个,新增元素应该进入小根堆
            max_heap.push(num);
            int temp = max_heap.top();
            max_heap.pop();

            // 但不能直接把num 放进去,而是把含num的大根堆中的最小值放进去
            min_heap.push(temp);
        } else {
            // 奇数的逻辑同上
            min_heap.push(num);
            
            int temp = min_heap.top();
            min_heap.pop();
            max_heap.push(temp);
        }
        
        isEven = !isEven;
    }

    double GetMedian() { 
        // 如果是偶数,说明两个堆相等,取两个堆顶的平均数
        if (isEven) {
            return (min_heap.top() + max_heap.top()) * 0.5;
        } else {
            // 如果是奇数,根据前面的定义,小根堆比大根堆多一个,所以是小根堆的堆顶
            return min_heap.top() * 1.0;
        }
    }

};
  • 时间复杂度: Insert的时间复杂度是O(logn)O(logn),因为只涉及了堆的调整操作。GetMedianO(1)O(1)
  • 空间复杂度:O(n)O(n)