【劍指offer】資料流中的中位數(最大最小堆實現)
阿新 • • 發佈:2018-11-29
題目描述
如何得到一個數據流中的中位數?如果從資料流中讀出奇數個數值,那麼中位數就是所有數值排序之後位於中間的數值。如果從資料流中讀出偶數個數值,那麼中位數就是所有數值排序之後中間兩個數的平均值。我們使用Insert()方法讀取資料流,使用GetMedian()方法獲取當前讀取資料的中位數。
構建一個最大堆和最小堆,並且保證最大堆中的資料都要小於最小堆的資料,也就是最大堆的根節點小於最小堆的根節點,插入的時候交替插入
如果插入到最大堆:如果輸入小於最大堆的根節點,插入到最大堆,如果大於最大堆的根節點,插入最小堆,將最小堆的根節點插入到左邊的最大堆
如果插入到最小堆:如果輸入的大於最小堆的根節點,插入到最小堆,如果小於最小堆的根節點,那麼插入最大堆,最大堆的根節點插入到右邊的最小堆
import heapq
class Solution:
def __init__(self):
self.maxHeap = []
self.minHeap = []
self.flag = 1
def Insert(self, num):
# write code here
if self.flag:
if self.minHeap == [] or num <= self.minHeap[0]: #插入最大堆
self.maxHeap. append(num) #插入
heapq.heapify(self.maxHeap) #調整
elif num > self.minHeap[0]: #大於最小堆最小元素,需要插入到最小堆,並且替換
self.maxHeap.append(heapq.heappop(self.minHeap))
heapq._heapify_max(self.maxHeap)
heapq.heappush(self.minHeap, num)
self. flag = 0
else:
if self.maxHeap == [] or num >= self.maxHeap[0]: #插入最小堆
heapq.heappush(self.minHeap, num)
elif num < self.maxHeap[0]: #插入最大堆,並把最大堆的根替換到右邊的最小堆
heapq.heappush(self.minHeap, heapq.heappop(self.maxHeap)) #最大堆的根
self.maxHeap.append(num)
heapq._heapify_max(self.maxHeap)
self.flag = 1
def GetMedian(self):
# write code here
if self.flag:
return float('{:.2f}'.format((self.maxHeap[0] + self.minHeap[0]) / 2))
else:
return float('{:.2f}'.format(self.maxHeap[0]))
測試
sol = Solution()
sol.Insert(5)
sol.Insert(2)
sol.Insert(6)
sol.Insert(7)
sol.Insert(8)
print(sol.maxHeap)
print(sol.minHeap)
print(sol.GetMedian())
輸出:
[6, 2, 5]
[7, 8]
6.0