LeetCode 295. Find Median from Data Stream 题解

题目描述

设计一个数据结构，支持一下两种操作：

• void addNum(int num) - 从数据流中增加一个数到该数据结构
• double findMedian() - 返回目前为止数据流的中位数

解题思路

单纯实现该数据结构的功能非常简单，最 Naive 的思路，用 list 存储数据，排序即可找出中位数，该方法的时间复杂度是 O(N*log(N))。但是有更 efficient 的方法，用两个堆实现 O(log(N)) 的复杂度。
用小堆samll存储当前数据中较小的一半，大堆large存储当前数据中较大一半，其中small的长度不大于large的长度。计算中位数时只需取堆的栈顶元素即可，时间复杂度为 O(1)，向堆中增加数字的时间复杂度为 O(log(N))。

代码实现

class MedianFinder(object):

    def __init__(self):
        """
        initialize your data structure here.
        """

        self.small = []
        self.large = []

    def addNum(self, num):
        """
        :type num: int
        :rtype: None
        """

        if len(self.small) == len(self.large):
            heapq.heappush(self.large, -heapq.heappushpop(self.small, -num))
        else:
            heapq.heappush(self.small, -heapq.heappushpop(self.large, num))

    def findMedian(self):
        """
        :rtype: float
        """

        if len(self.small) == len(self.large):
            return (-self.small[0] + self.large[0]) / 2.0
        else:
            return self.large[0]