[剑指 Offer 第 2 版第 41 题] “数据流中的中位数”做题记录

传送门：牛客网 online judge 地址。

我们这里使用 AcWing 上面的在线测评系统，传送门：《剑指 Offer 》（第 2 版）第 41 题：数据流中的中位数。

题目描述

如何得到一个数据流中的中位数？

如果从数据流中读出奇数个数值，那么中位数就是所有数值排序之后位于中间的数值。

如果从数据流中读出偶数个数值，那么中位数就是所有数值排序之后中间两个数的平均值。

样例

  输入：1, 2, 3, 4

    输出：1,1.5,2,2.5

    解释：每当数据流读入一个数据，就进行一次判断并输出当前的中位数。

求解思路与关键

同 LeetCode 第 295 题，传送门：数据流的中位数。

思路：借助两个堆，一个小顶堆，一个大顶堆。

任何时刻，两个堆中应该始终保持的性质：

1、小顶堆中最小的元素，都不会小于大顶堆中的元素；

2、小顶堆的元素数目或者与大顶堆的元素数组相等，或者多 1 。

因此，我们可以针对当前元素的个数进行分类讨论：

1、当前元素个数为偶数的时候，比如一开始的时候，读入一个元素（最终应该是小顶堆中多一个元素），此时元素个数是奇数，我们应该返回小顶堆中的堆顶元素；

如何保持性质：因为最终应该是小顶堆中多一个元素，所以先经过大顶堆，然后再到小顶堆。

2、当前元素个数为奇数的时候，此时再读入一个元素（最终应该是大顶堆中多一个元素），元素个数是偶数，我们应该返回大顶堆中的堆顶元素和小顶堆中的堆顶元素的平均数。

如何保持性质：因为最终应该是大顶堆中多一个元素，所以先经过小顶堆，然后再到大顶堆。

注意：Python 中的 heap 只有小顶堆，在构造大顶堆的时候，要绕一个弯子，把相反数 push 进去，pop 出来的时候也要取相反数。

参考解答

参考解答1

Python 代码：

  import heapq


    class Solution:

        def __init__(self):
            self.min_heap = []
            self.max_heap = []
            self.count = 0

        def insert(self, num):
            """
            :type num: int
            :rtype: void
            """
            self.count += 1
            if self.count % 2 == 1:
                heapq.heappush(self.max_heap, -num)
                temp = -heapq.heappop(self.max_heap)
                heapq.heappush(self.min_heap, temp)
            else:
                heapq.heappush(self.min_heap, num)
                temp = heapq.heappop(self.min_heap)
                heapq.heappush(self.max_heap, -temp)

            # print(self.min_heap)
            # print(self.max_heap)

        def getMedian(self):
            """
            :rtype: float
            """
            if self.count % 2 == 1:
                return self.min_heap[0]
            else:
                return (self.min_heap[0] + (-self.max_heap[0])) / 2


    if __name__ == '__main__':
        solution = Solution()
        solution.insert(1)
        result = solution.getMedian()
        print(result)
        solution.insert(2)
        result = solution.getMedian()
        print(result)

        solution.insert(3)
        result = solution.getMedian()
        print(result)

        solution.insert(4)
        result = solution.getMedian()
        print(result)

说明：1、奇数偶数的判断可以使用位运算：count & 1 == 1 可以判断 count 是不是奇数；

2、上面的代码写得有一些烦琐，我们注意到，反正小顶堆中的元素一定 >= 大顶堆中的元素，我们可以把顺序统一一下，如果当前是奇数的话，新来的一个数，先进入小顶堆，然后小顶堆中出一个数再进入大顶堆，这样小顶堆和大顶堆元素数目就一样了。那如果当前是偶数，最终效果是小顶堆多一个元素，还按照刚刚那个流程，我们从大顶堆再拿一个元素放回小顶堆就可以了。

我们统一写在参考解答2里。

参考解答2：

Python 代码

  import heapq


    class Solution:

        def __init__(self):
            self.min_heap = []
            self.max_heap = []
            self.count = 0

        def insert(self, num):
            """
            :type num: int
            :rtype: void
            """
            # 当前是奇数的时候，直接"最小堆" -> "最大堆"，就可以了
            # 此时"最小堆" 与 "最大堆" 的元素数组是相等的

            # 当前是偶数的时候，"最小堆" -> "最大堆"以后，最终我们要让"最小堆"多一个元素
            # 所以应该让 "最大堆" 拿出一个元素给 "最小堆"

            heapq.heappush(self.min_heap, num)
            temp = heapq.heappop(self.min_heap)
            heapq.heappush(self.max_heap, -temp)
            if self.count & 1 == 0:
                temp = -heapq.heappop(self.max_heap)
                heapq.heappush(self.min_heap, temp)
            self.count += 1
            # print(self.min_heap)
            # print(self.max_heap)

        def getMedian(self):
            """
            :rtype: float
            """
            if self.count & 1 == 1:
                return self.min_heap[0]
            else:
                return (self.min_heap[0] + (-self.max_heap[0])) / 2

作者：liweiwei1419

来源：https://liweiwei1419.github.io/sword-for-offer/