程序员面试金典 - 面试题 17.20. 连续中值

最新推荐文章于 2024-07-18 22:43:28 发布

随便发挥

最新推荐文章于 2024-07-18 22:43:28 发布

阅读量290

点赞数

分类专栏： Leetcode 文章标签：面试算法数据结构

本文链接：https://blog.csdn.net/zjulyx1993/article/details/126435956

版权

Leetcode 专栏收录该内容

309 篇文章 4 订阅

订阅专栏

题目难度: 困难

原题链接

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

题目描述

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

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

例如，

[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

题目思考

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

解决方案

思路

分析题目, 一个很容易想到的方案就是手动维护有序数组, 然后插入时二分查找插入位置, 查询时使用中间下标
不过这样虽然查询复杂度是 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]