代码多,建议电脑观看或者点击原文阅读
今天来看两道使用两个堆在变化的序列中取中位数的问题~
此题[295. Find Median from Data Stream ](https://leetcode.com/problems/find-median-from-data-stream/description/) 与[480. Sliding Window Median ](https://leetcode.com/problems/sliding-window-median/) 类似。都是在一个变化的序列中多次求中位数,而我们知道中位数与大小相关,所以很容易想到用堆来维护。
但对于堆来说,我们只能方便的取最大或最小,在这个问题中,我们可以用一个最小堆和一个最大堆来维护有序性,并且维持两个堆内元素个数之差始终在1以内。那么每次取中位数就只需要O(1)的时间,每次插入元素只需要O(n)的时间。具体如下:
初始化max_heap=[]用来保存小于中位数的数, min_heap=[]用来保存大于或等于中位数的数
插入一个元素时
根据大小判断应该插入到max_heap还是min_heap中
通过heappop和heappush保证len(max_heap)
取中位数时
如果len(min_heap)==len(max_heap)+1, 取min_heap堆顶元素
如果len(min_heap)==len(max_heap), 取两堆堆顶元素平均数
代码如下:
fromheapqimport*
classMedianFinder:
def__init__(self):
self.heaps= [], []
defaddNum(self,num):
small,large=self.heaps
heappush(small,-heappushpop(large,num))
iflen(large)
heappush(large,-heappop(small))
deffindMedian(self):
small,large=self.heaps
iflen(large)>len(small):
returnfloat(large[])
return(large[]-small[])/2.0
如果是[480. Sliding Window Median ](https://leetcode.com/problems/sliding-window-median/) ,只需要在addNum时增加判断是否超出窗口大小,如果超过,就invalid离开窗口的元素,否则直接加入。
代码如下(有点啰嗦,没来得及优化):
classSolution(object):
defmedianSlidingWindow(self,nums,k):
"""
:type nums: List[int]
:type k: int
:rtype: List[float]
"""
ifk==1:
returnlist(map(float,nums))
defpop(heap):
clean(heap)
returnheapq.heappop(heap)
deftop(heap):
clean(heap)
returnheap[]
defclean(heap):
whileheapandnotheap[][1]:
heapq.heappop(heap)
defmed(heap1,heap2):
clean(heap1)
clean(heap2)
ifk&1:
returnheap2[][]
else:
return(heap2[][]-heap1[][])/2
N=len(nums)
nums= [[nums[i],True,1]foriinrange(len(nums))]
heap1= []
heap2= []
foriinrange(k):
nums[i][] =-nums[i][]
heapq.heappush(heap1,nums[i])
foriinrange((k+1)//2):
n=pop(heap1)
n[] =-n[]
n[2] =2
heapq.heappush(heap2,n)
ret= [med(heap1,heap2)]
foriinrange(N-k):
nums[i][1] =False
ifnums[i][2]==1:
ifnums[i+k][]
nums[i+k][] =-nums[i+k][]
heapq.heappush(heap1,nums[i+k])
ret.append(med(heap1,heap2))
else:
nums[i+k][2] =2
heapq.heappush(heap2,nums[i+k])
n=pop(heap2)
n[2] =1
n[] =-n[]
heapq.heappush(heap1,n)
ret.append(med(heap1,heap2))
else:
ifnums[i+k][]>=-top(heap1)[]:
nums[i+k][2] =2
heapq.heappush(heap2,nums[i+k])
ret.append(med(heap1,heap2))
else:
nums[i+k][] =-nums[i+k][]
heapq.heappush(heap1,nums[i+k])
n=pop(heap1)
n[2] =2
n[] =-n[]
heapq.heappush(heap2,n)
ret.append(med(heap1,heap2))
returnlist(map(float,ret))
领取专属 10元无门槛券
私享最新 技术干货