算法题:数据流的中位数
题目分析
nums = [1, 5, 3, 4, 2]
,那么在流的每个时刻,我们都要能快速计算出当前中位数。每次接收一个新的数字; 求出当前所有数字的中位数; 要求这个操作尽可能高效,不然你就得不停排序,效率堪忧。
直觉法:排序求中位数
sort()
来搞定。我们每次接收到一个新的数字,都把当前数字加入数组,然后排序数组,最后取数组的中位数。import bisect
class MedianFinder:
def __init__(self):
self.nums = []
def addNum(self, num: int) -> None:
bisect.insort(self.nums, num)
def findMedian(self) -> float:
n = len(self.nums)
if n % 2 == 1:
return self.nums[n // 2]
else:
return (self.nums[n // 2 - 1] + self.nums[n // 2]) / 2.0
addNum
方法通过 bisect.insort
将新的数字插入到已有列表中,保证了列表始终是有序的。然后 findMedian
方法直接根据列表的长度来判断中位数是列表的中间元素还是中间两个元素的平均值。这里的 bisect.insort
是一个二分查找插入方法,时间复杂度为 O(log n),比直接用 sort()
排序要高效一些。高效解法:堆结构
最大堆的堆顶元素是当前小部分中的最大值; 最小堆的堆顶元素是当前大部分中的最小值; 中位数要么是最大堆的堆顶元素(当元素个数是奇数时),要么是最大堆和最小堆堆顶元素的平均值(当元素个数是偶数时)。
heapq
模块来实现时,heapq
默认是最小堆,因此我们需要把最大堆模拟成负数的最小堆。import heapq
class MedianFinder:
def __init__(self):
# 最大堆(小部分)
self.max_heap = []
# 最小堆(大部分)
self.min_heap = []
def addNum(self, num: int) -> None:
# 先把数字放到合适的堆里
if len(self.max_heap) == 0 or num <= -self.max_heap[0]:
heapq.heappush(self.max_heap, -num) # 使用负数模拟最大堆
else:
heapq.heappush(self.min_heap, num)
# 保证最大堆和最小堆的元素个数平衡
if len(self.max_heap) > len(self.min_heap) + 1:
heapq.heappush(self.min_heap, -heapq.heappop(self.max_heap))
if len(self.min_heap) > len(self.max_heap):
heapq.heappush(self.max_heap, -heapq.heappop(self.min_heap))
def findMedian(self) -> float:
if len(self.max_heap) > len(self.min_heap):
return -self.max_heap[0]
else:
return (-self.max_heap[0] + self.min_heap[0]) / 2.0
讲解代码
初始化:我们用两个堆来维护数据流的两部分, max_heap
存储较小的一半数据(其实存的是负数,模拟成最大堆),min_heap
存储较大的一半数据。添加数字:每次添加数字时,我们根据数字的大小决定放入哪个堆。如果数字小于等于最大堆的堆顶(记得是负数),就放到最大堆;否则就放到最小堆。 平衡堆:每次插入完数字后,我们要确保两个堆的大小平衡。最大堆的大小不能超过最小堆太多,最小堆的大小不能超过最大堆。 查找中位数:当最大堆的元素个数比最小堆多时,中位数就是最大堆的堆顶元素;如果两个堆的大小相等,中位数就是两个堆顶元素的平均值。
总结
对编程、职场感兴趣的同学,大家可以联系我微信:golang404,拉你进入“程序员交流群”。
虎哥作为一名老码农,整理了全网最全《python高级架构师资料合集》。