设计 | 数据流中的中位数

题目

求数据流的中位数

• 设计一个类，包含addNum和findMedian两个方法，findMedian时间复杂度优化到O(1)

方法

1. 最简单的方法 - 排序，每次addNum时，元素追加到list末尾；findMedian时快速排序，然后求得中位数；addNum和findMedian的时间复杂度分别是O(n) 和 O(nlogn)；提供这个答案明显还不够；如何优化findMedian的时间复杂度呢？
2. 维护有序list，二分查找中位数；addNum时保持list有序，二分查找找到小于num的最小元素的索引idx，然后插入新元素到list；addNum和findMedian时间复杂度分别为O(n) O(logn)
3. 维护两个heap；求中位数，没必要保持所有元素排序；只需要将元素等分为两个堆heap1，heap2，保证all num in heap1 <= all num in heap2；addNum和findMedian时间复杂度分别为O(logn) O(1)

class MedianFinder:
    def __init__(self):
        """
        initialize your data structure here.
        """
        self.heap1 = []
        self.heap2 = []

    def addNum(self, num: int) -> None:
        l1, l2 = len(self.heap1), len(self.heap2)
        if l2 == 0 or self.heap2[0] <= num:
            heapq.heappush(self.heap2, num)
        else:
            heapq.heappush(self.heap1, -num)
        l1, l2 = len(self.heap1), len(self.heap2)      
        if l2 - l1 > 1:
            item = heapq.heappop(self.heap2)
            heapq.heappush(self.heap1, -item)
        elif l1 > l2:
            item = heapq.heappop(self.heap1)
            heapq.heappush(self.heap2, -item)
        # print(self.heap1, self.heap2)

    def findMedian(self) -> float:
        l1, l2 = len(self.heap1), len(self.heap2)
        if l1 + l2 == 0:
            return 0
        if (l1 + l2) % 2 == 1:
            return self.heap2[0]
        else:
            return (-self.heap1[0] + self.heap2[0]) / 2


# Your MedianFinder object will be instantiated and called as such:
# obj = MedianFinder()
# obj.addNum(num)
# param_2 = obj.findMedian()

• 这个题目前做到了一遍过