Leetcode 295. Find Median from Data Stream (python+cpp)

题目

解法1：simple sorting

每次找中位数的操作都将列表进行重新排序，这种解法时间复杂度为O(nlogn)

class MedianFinder:def __init__(self):"""initialize your data structure here."""self.numList = []def addNum(self, num: int) -> None:self.numList.append(num)def findMedian(self) -> float:self.numList = sorted(self.numList)n = len(self.numList)if n%2 != 0:return self.numList[n//2]else:return (self.numList[n//2 -1] + self.numList[n//2])/2

解法2：插入

想办法保持列表是一个有序的序列，这样就可以做到每次查询中位数只需要O(1)的复杂度。保持增序可以通过在添加新元素时以插入的方式来进行。插入最简单的就是每次从一个方向去找，然后直到找到合适的位置，这种时间复杂度为O(n)。很奇怪的是这种复杂度比上面的要小，但是上面的可以通过，这种解法却TLE。

class MedianFinder:def __init__(self):"""initialize your data structure here."""self.numList = []def addNum(self, num: int) -> None:if not self.numList:self.numList.append(num)else:for i in range(len(self.numList)):if self.numList[i]>num:self.numList = self.numList[:i] + [num] +self.numList[i:]breakelif i==len(self.numList)-1:self.numList = self.numList + [num]breakdef findMedian(self) -> float:n = len(self.numList)if n%2 != 0:return self.numList[n//2]else:return (self.numList[n//2 -1] + self.numList[n//2])/2

当然这边可能的优化是将插入的过程通过binary search来实现

解法3：two heaps

构建两个堆，一个最小堆，一个最大堆，永远保持这两个堆的元素个数差1.这样的话中位数就可以通过两个堆顶元素来计算。具体流程如下：

• 首先由于python中只有最小堆没有最大堆，所以只能对元素取相反数来模拟最大堆
• 如果两个堆都为空，随便将元素加入其中一个堆
• 将现在要插入的元素与两个堆中的一个的顶部元素进行比较，然后根据比较结果将元素插入到其中一个堆。比如如果当前元素大于最小堆的堆顶元素，那么这个元素一定是插入最小堆的，否则插入最大堆
• 最关键的是在每次插入之后，对两个堆进行pop和push操作，保证这两个堆的元素个数不大于1
  这样的解法只会消耗log(n)的时间复杂度

class MedianFinder:def __init__(self):"""initialize your data structure here."""minheap,maxheap = [],[]heapq.heapify(minheap)heapq.heapify(maxheap)self.minheap = minheapself.maxheap = maxheapdef addNum(self, num: int) -> None:if not self.minheap and not self.maxheap:heapq.heappush(self.maxheap,-num)elif self.maxheap and num < -self.maxheap[0]:heapq.heappush(self.maxheap,-num)else:heapq.heappush(self.minheap,num)self.rebalance()def rebalance(self):len_diff = len(self.maxheap) - len(self.minheap)if len_diff>1:tmp = heapq.heappop(self.maxheap)heapq.heappush(self.minheap,-tmp)elif len_diff<-1:tmp = heapq.heappop(self.minheap)heapq.heappush(self.maxheap,-tmp)def findMedian(self) -> float:if not self.maxheap and not self.minheap:return 0if len(self.maxheap)>len(self.minheap):return -self.maxheap[0]elif len(self.minheap)>len(self.maxheap):return self.minheap[0]else:return (-self.maxheap[0]+self.minheap[0])/2

参考链接：https://leetcode.com/problems/find-median-from-data-stream/discuss/806419/Python-MinMax-heaps-(with-comments)

C++版本

class MedianFinder {
    
public:priority_queue<int> max_heap;priority_queue<int, vector<int>, greater<int>> min_heap;/** initialize your data structure here. */MedianFinder() {
    }void addNum(int num) {
    if(min_heap.empty() && max_heap.empty()){
    max_heap.push(num);}else if(!min_heap.empty() && num > min_heap.top()){
    min_heap.push(num);}else{
    max_heap.push(num);}rebalance();}void rebalance(){
    int len_diff = min_heap.size() - max_heap.size();int tmp;if(len_diff > 1){
    tmp = min_heap.top();min_heap.pop();max_heap.push(tmp);}if(len_diff < -1){
    tmp = max_heap.top();max_heap.pop();min_heap.push(tmp);}}double findMedian() {
    // cout << min_heap.size() << max_heap.size() << endl;if(min_heap.empty() && max_heap.empty()){
    return 0;}else if(min_heap.size() > max_heap.size()){
    return min_heap.top();}else if(min_heap.size() < max_heap.size()){
    return max_heap.top();}else{
    return (min_heap.top() + max_heap.top())/2.0;}}
};