295. Find Median from Data Stream

重温学会双堆法的经典题目

题干：

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.

解释：

有一个输入流，要求返回任意时刻的中位数

思考：

比较常规的做法是每次插入一个数，都直接找到插入位置然后保持顺序插入。但是插入的时间复杂度是On，如果有n次操作，时间复杂度为N2. 所以这时候就要从中位数的定义上想办法，其实中位数就等于小于他的数的最大值，和大于他的数的最小值，所以我们用两个堆来存储左边和右边，然后调整堆的大小使得平衡即可。

答案：

import heapq
class MedianFinder(object):
    
    def __init__(self):
        """
        initialize your data structure here.
        """
        #最大堆
        self.small =[]
        #总是使得最大堆的长度《=最小堆堆长度+1
        #最小堆
        self.big   =[]
        

    def addNum(self, num):
        """
        :type num: int
        :rtype: None
        """
        heapq.heappush(self.big,num)
        self.balance(self.small,self.big)
        
    def balance(self,small,big):
        if (len(small)+len(big))%2 ==1:
            while len(small)<len(big)+1:
                heapq.heappush(small,-heapq.heappop(big))
        else:
            while len(small)!=len(big):
                heapq.heappush(small,-heapq.heappop(big))

    def findMedian(self):
        """
        :rtype: float
        """
        if len(self.small) == len(self.big):
            return float((-self.small[0]+self.big[0])/2)
        else :
            return -self.small[0]
def main():
    sol = MedianFinder()
    sol.findMedian()
    sol.addNum(1)
    sol.addNum(2)
    print(sol.findMedian())
main()

答案补充：