如何使用大小堆解决数据流中位数问题

2024-02-18 21:58:41

算法流程：

初始化两个堆：大顶堆和小顶堆。大顶堆保存的是数据流中较大的部分，小顶堆保存的是数据流中较小的部分。
将数据流中的第一个数添加到大顶堆中。
将数据流中的第二个数添加到小顶堆中。
判断大顶堆和小顶堆的元素个数是否相等。
如果大顶堆和小顶堆的元素个数相等，则数据流中的中位数是这两个堆的堆顶元素的平均值。
如果大顶堆的元素个数多于小顶堆的元素个数，则数据流中的中位数是大顶堆的堆顶元素。
如果小顶堆的元素个数多于大顶堆的元素个数，则数据流中的中位数是小顶堆的堆顶元素。
将数据流中的下一个数添加到对应堆中，并重复步骤4-7。

时间复杂度：
该算法的时间复杂度是 O(log N)，因为在每次添加一个数到堆中时，只需要对堆进行一次调整，而调整堆的时间复杂度是 O(log N)。

空间复杂度：
该算法的空间复杂度是 O(N)，因为需要存储数据流中的所有数。

代码实现：

import heapq

class MedianFinder:
    def __init__(self):
        self.max_heap = []  # 大顶堆
        self.min_heap = []  # 小顶堆

    def add_num(self, num):
        if len(self.max_heap) == len(self.min_heap):
            heapq.heappush(self.max_heap, -heapq.heappushpop(self.min_heap, num))
        else:
            heapq.heappush(self.min_heap, heapq.heappushpop(self.max_heap, -num))

    def find_median(self):
        if len(self.max_heap) == len(self.min_heap):
            return (self.min_heap[0] - self.max_heap[0]) / 2
        else:
            return self.max_heap[0] * -1


if __name__ == '__main__':
    mf = MedianFinder()
    mf.add_num(1)
    mf.add_num(2)
    mf.add_num(3)
    print(mf.find_median())  # 2.0
    mf.add_num(4)
    print(mf.find_median())  # 2.5