参考:https://blog.csdn.net/xinkengruo3162/article/details/88928408
随时返回数据流中的中位数
小数放大根堆,大数放小根堆,保持两个堆的平衡,最后取出两个堆的头节点即可
import heapq
#实现大根堆,heapy默认的是小根堆
class bigHeap():
def __init__(self):
self.arr=list()
def heap_insert(self,val):
heapq.heappush(self.arr,-val) #加入相反数,实现大根堆
def heapify(self):
heapq.heapify(self.arr)
def heap_pop(self): #弹出时也要取反
return -heapq.heappop(self.arr)
def get_top(self): 返回堆顶元素,也要取反
if not self.arr:
return
return -self.arr[0]
#实现小根堆
class smallHeap():
def __init__(self):
self.arr=list()
def heap_insert(self,val):
heapq.heappush(self.arr,val)
def heapify(self):
heapq.heapify(self.arr)
def heap_pop(self):
return heapq.heappop(self.arr)
def get_top(self):
if not self.arr:
return
return self.arr[0]
#随时返回数据流中的中位数
class medianHolder():
def __init__(self):
self.bigheap=bigHeap()
self.smallheap=smallHeap()
def addnum(self,num): #加数操作
if len(self.bigheap.arr)==0:
self.bigheap.heap_insert(num)
return
if self.bigheap.get_top()>=num:
self.bigheap.heap_insert(num)
else:
if len(self.smallheap.arr) == 0:
self.smallheap.heap_insert(num)
return
if self.smallheap.get_top()>num:
self.bigheap.heap_insert(num)
else:
self.smallheap.heap_insert(num)
self.modifyTwoheap()
def modifyTwoheap(self): #调整两个堆的大小
smallszize=len(self.smallheap.arr)
bigsize=len(self.bigheap.arr)
if smallszize==bigsize+2:
self.bigheap.heap_insert(self.smallheap.heap_pop())
if bigsize==smallszize+2:
self.smallheap.heap_insert(self.bigheap.heap_pop())
def getmedian(self): #获取中位数
smallszize = len(self.smallheap.arr)
bigsize = len(self.bigheap.arr)
if smallszize + bigsize == 0:
return None
smallhead=self.smallheap.get_top()
bighead=self.bigheap.get_top()
if (smallszize + bigsize) %2 == 0:
return (smallhead+bighead)/2
else:
return smallhead if smallszize>bigsize else bighead
if __name__=="__main__":
arr=[68,51,42,92,13,46,24,58,62,72,89]
print(sorted(arr))
median=medianHolder()
for i in range(len(arr)):
median.addnum(arr[i])
print(median.getmedian())
返回最大的k个数
用大根堆来实现
import heapq
class bigHeap():
def __init__(self):
self.arr=list()
def heap_insert(self,val):
heapq.heappush(self.arr,-val)
def heapify(self):
heapq.heapify(self.arr)
def heap_pop(self):
return -heapq.heappop(self.arr)
def get_top(self):
if not self.arr:
return
return -self.arr[0]
if __name__=="__main__":
arr=[68,51,42,92,13,46,24,58,62,72,89]
print(sorted(arr))
k=3
s=bigHeap()
for i in range(len(arr)):
s.heap_insert(arr[i])
for i in range(k):
print(s.heap_pop(),end=" ")