方法一:有序插入
from sortedcontainers import SortedList
class MedianFinder:
def __init__(self):
#self.nums = []
self.nums = SortedList()
def addNum(self, num: int) -> None:
# 插入后排序
#self.nums.append(num)
#self.nums.sort()
# 使用 bisect 有序插入
#bisect.insort(self.nums,num)
# 二分有序插入
#n = len(self.nums)
#l, r = 0, n
# while l < r:
# mid = (l + r)//2
# if num > self.nums[mid]:
# l = mid + 1
# else:
# r = mid
#self.nums.insert(r,num)
# 使用 SortedList
self.nums.add(num)
def findMedian(self) -> float:
# n = len(self.nums)
return self.nums[n//2] if (n:=len(self.nums))%2 else (self.nums[n//2]+self.nums[n//2-1]) / 2
方法二:对顶堆
大根堆维护比中位数小的数,小根堆维护比中位数大的数。
若 num 比中位数大(小根堆堆顶元素),则将其插入小根堆,插入后若小根堆长度大于大根堆长度,则将其堆顶转移到大根堆;否则将其插入大根堆。保证大根堆长度比小根堆长度多1或者相等。中位数要么在大根堆堆顶,要么就是取两个堆的堆顶求平均。
from sortedcontainers import SortedList
class MedianFinder:
def __init__(self):
"""
initialize your data structure here.
"""
self.lo = [] # 大根堆
self.hi = [] # 小根堆
def addNum(self, num: int) -> None:
heapq.heappush(self.lo, -num)
heapq.heappush(self.hi, -heapq.heappop(self.lo))
if len(self.lo) < len(self.hi):
heapq.heappush(self.lo, -heapq.heappop(self.hi))
def findMedian(self) -> float:
return -self.lo[0] if len(self.lo) > len(self.hi) else (-self.lo[0] + self.hi[0]) / 2
class MedianFinder:
def __init__(self):
self.left = []
self.right = []
def addNum(self, num: int) -> None:
l, r, = self.left, self.right
if not l or num <= -l[0]: # left 为空或 num <= 中位数(大根堆堆顶元素)
heapq.heappush(l, -num)
if len(l) - len(r) > 1:
heapq.heappush(r, -heapq.heappop(l)) # 转移小堆顶到大堆
else:
heapq.heappush(r, num)
if len(r) > len(l):
heapq.heappush(l, -heapq.heappop(r))
def findMedian(self) -> float:
return -self.left[0] if len(self.left) > len(self.right) else (-self.left[0]+self.right[0])/2