[295][困难][二分][堆] 数据流的中位数

解题思路

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

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

例如，

[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

进阶:

• 如果数据流中所有整数都在 0 到 100 范围内，你将如何优化你的算法？

• 如果数据流中 99% 的整数都在 0 到 100 范围内，你将如何优化你的算法？

解题思路

二分插入排序

维护一个列表, 使这个列表中的元素保持递增的状态. 对于数据流新来的一个数字, 使用二分法决定插入的位置, 更新数组. 在计算当前中位数时根据数组长度的奇偶性进行调整.

class MedianFinder:

    def __init__(self):
        """
        initialize your data structure here.
        """
        self.nums = []

    def addNum(self, num: int) -> None:
        index = bisect.bisect(self.nums, num)
        self.nums.insert(index, num)

    def findMedian(self) -> float:
        if len(self.nums) == 0:
            return None
        middle = len(self.nums) // 2
        return (self.nums[middle - 1] + self.nums[middle]) / 2 if len(self.nums) % 2 == 0 else self.nums[middle]


# Your MedianFinder object will be instantiated and called as such:
# obj = MedianFinder()
# obj.addNum(num)
# param_2 = obj.findMedian()

两个堆

面试题41. 数据流中的中位数（优先队列 / 堆，清晰图解）

import heapq


class MedianFinder:

    def __init__(self):
        """
        initialize your data structure here.
        """
        self.larger = []
        self.smaller = []

    def addNum(self, num: int) -> None:
        if len(self.larger) == len(self.smaller):
            heapq.heappush(self.smaller, -num)
            heapq.heappush(self.larger, -heapq.heappop(self.smaller))
        else:
            heapq.heappush(self.larger, num)
            heapq.heappush(self.smaller, -heapq.heappop(self.larger))

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

优化下细节, 在选定要更新哪个堆之后, 比较当前数字与另一个堆顶数字大小, 可能可以免去另一个堆寻找最值的过程:

import heapq


class MedianFinder:

    def __init__(self):
        """
        initialize your data structure here.
        """
        self.larger = []
        self.smaller = []

    def addNum(self, num: int) -> None:
        if len(self.larger) == len(self.smaller):
            if len(self.smaller) and num >= -self.smaller[0]:
                heapq.heappush(self.larger, num)
            else:
                heapq.heappush(self.smaller, -num)
                heapq.heappush(self.larger, -heapq.heappop(self.smaller))
        else:
            if len(self.larger) and num <= self.larger[0]:
                heapq.heappush(self.smaller, -num)
            else:
                heapq.heappush(self.larger, num)
                heapq.heappush(self.smaller, -heapq.heappop(self.larger))

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