题目难度: 困难

****

今天继续更新程序员面试金典系列, 大家在公众号 算法精选 里回复 面试金典 就能看到该系列当前连载的所有文章了, 记得关注哦~

题目描述

随机产生数字并传递给一个方法。你能否完成这个方法，在每次产生新值时，寻找当前所有值的中间值（中位数）并保存。

中位数是有序列表中间的数。如果列表长度是偶数，中位数则是中间两个数的平均值。

例如，

[2,3,4] 的中位数是 3

[2,3] 的中位数是 (2 + 3) / 2 = 2.5

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

  • void addNum(int num) - 从数据流中添加一个整数到数据结构中。
  • double findMedian() - 返回目前所有元素的中位数。

示例：

  • addNum(1)
  • addNum(2)
  • findMedian() -> 1.5
  • addNum(3)
  • findMedian() -> 2

题目思考

  1. 如何尽可能优化时间复杂度?

解决方案

思路

  • 分析题目, 一个很容易想到的方案就是手动维护有序数组, 然后插入时二分查找插入位置, 查询时使用中间下标
  • 不过这样虽然查询复杂度是 O(1), 但插入复杂度达到了 O(N), 效率较低, 有没有更优的方案呢?
  • 根据中位数性质, 它要么是偶数长度数组的两个中间值的平均数, 要么是奇数长度数组的最中间的值
  • 注意到中间值左边的部分一定是小于中间值的, 而右边的部分一定是大于中间值的, 是有一定的有序性的
  • 我们可以利用这一点, 构造两个堆:
    • 左边是一个大顶堆, 存放所有小于等于中间值的数(奇数长度的话, 堆顶就是中间值)
    • 右边是一个小顶堆, 存放所有大于等于中间值的数(因为可能有很多重复元素)
  • 查询中位数
    • 当前元素个数为奇数时, 直接返回左堆顶
    • 当前元素个数为偶数时, 返回左堆顶和右堆顶的平均值
  • 插入新元素
    • 如果当前左堆和右堆个数相同, 那么左堆需要增加一个元素
      • 具体是增加新元素还是右堆顶, 则要看当前 num 和右堆顶的大小
      • 当前 num 更小的话直接插入左边即可
      • 否则就要先把右堆顶先挪到左堆, 然后右堆再插入新元素
    • 如果当前左堆比右堆多 1, 那么右堆需要增加一个元素
      • 具体是增加新元素还是左堆顶, 则要看当前 num 和左堆顶的大小
      • 当前 num 更大的话直接插入右边即可
      • 否则就要先把左堆顶先挪到右堆, 然后左堆再插入新元素
  • 下面代码有详细的注释, 方便大家理解

复杂度

  • 时间复杂度 O(logN): 查询和插入都只需要常数次堆操作, 复杂度为 O(logN)
  • 空间复杂度 O(N): 两个堆共需要存 N 个元素

代码

import heapq

class MedianFinder:
    def __init__(self):
        """
        initialize your data structure here.
        """
        # 使用两个堆, 左边大顶堆, 右边小顶堆
        # 注意python内置的heapq是小顶堆
        # 所以大顶堆的话需要将其值取相反数再插入, 也不要忘了pop的时候也要再取反回来
        self.left = []
        self.right = []

    def addNum(self, num: int) -> None:
        if len(self.left) == len(self.right):
            # 左堆需要增加一个元素
            if not self.left or num <= self.right[0]:
                # 左堆不存在, 或者新元素不大于右堆顶, 插到左边大顶堆
                heapq.heappush(self.left, -num)
            else:
                # 否则先把右堆顶插到左边, 然后右堆再插入新元素
                heapq.heappush(self.left, -heapq.heappop(self.right))
                heapq.heappush(self.right, num)
        else:
            # 右堆需要增加一个元素
            # 根据插入逻辑, 此时左堆至少有一个元素, 可以直接拿到左堆顶
            if num >= -self.left[0]:
                # 新元素不小于左堆顶, 直接插到右边小顶堆即可
                heapq.heappush(self.right, num)
            else:
                # 否则先把左堆顶插到右边, 然后左堆再插入新元素
                heapq.heappush(self.right, -heapq.heappop(self.left))
                heapq.heappush(self.left, -num)

    def findMedian(self) -> float:
        if len(self.left) == len(self.right):
            # 偶数个元素, 取两个堆顶的平均值
            return (-self.left[0] + self.right[0]) / 2
        else:
            # 奇数个元素, 左堆个数至少为1, 取左堆顶
            return -self.left[0]

大家可以在下面这些地方找到我~😊

我的 GitHub

我的 Leetcode

我的 CSDN

我的知乎专栏

我的头条号

我的牛客网博客

我的公众号: 算法精选, 欢迎大家扫码关注~😊

算法精选 - 微信扫一扫关注我