剑指 Offer 41. 数据流中的中位数

119 阅读2分钟

题目描述

如何得到一个数据流中的中位数?如果从数据流中读出奇数个数值,那么中位数就是所有数值排序之后位于中间的数值。如果从数据流中读出偶数个数值,那么中位数就是所有数值排序之后中间两个数的平均值。

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

设计一个支持以下两种操作的数据结构:
void addNum(int num) - 从数据流中添加一个整数到数据结构中。
double findMedian() - 返回目前所有元素的中位数。

示例 1:
输入:
["MedianFinder","addNum","addNum","findMedian","addNum","findMedian"]
[[],[1],[2],[],[3],[]]
输出:[null,null,null,1.50000,null,2.00000]

示例 2: 输入:
["MedianFinder","addNum","findMedian","addNum","findMedian"]
[[],[2],[],[3],[]]
输出:[null,null,2.00000,null,2.50000]

解题思路1: 暴力排序法 O(nlogn)

添加数据时, 直接将数据添加到一个数组, 获取中位数时, 将数组排序,获取中位数.

示例代码1: 略

解题思路2: 大顶堆+小顶堆

  1. 我们分别使用大顶堆和小顶堆来存储加入的数据, 大顶堆存储所有元素中较小的一部分, 小顶堆存储所有元素中较大的一部分. 这样分别遍历大顶堆和小顶堆后, 可以得到一个有序的数组.
  2. 获取中位数时: 当元素个数为奇数时, 直接拿小顶堆的堆顶元素就是中位数, 如果是偶数, 就拿2个堆顶元素计算中位数
  3. 添加元素时: 如果当前元素是偶数时, 为了保证可以直接拿小顶堆堆顶元素作为中位数, 所以我们要将一个元素加入到小顶堆, 那么加入哪个元素呢? 我们需要将当前元素先加入大顶堆, 然后将大顶堆的堆顶元素加入小顶堆. 如果当前元素是奇数时, 我们需要将当前元素先加入小顶堆, 然后将小顶堆的堆顶元素加入大顶堆.

Python中的堆默认是小顶堆, 所以我们将入堆,出堆的元素取反, 来实现一个大顶堆

示例代码2:

class MedianFinder:
    def __init__(self):
        self.small = []
        self.big = []

    def addNum(self, num: int) -> None:
        if len(self.small) != len(self.big):
            heapq.heappush(self.big, -heapq.heappushpop(self.small, num))
        else:
            heapq.heappush(self.small, -heapq.heappushpop(self.big, -num))
        print(self.small, self.big)

    def findMedian(self) -> float:
        return (
            self.small[0]
            if len(self.small) != len(self.big)
            else (self.small[0] - self.big[0]) / 2
        )