295. Find Median from Data Stream

Python

Two Heap

• 用兩個heap存大的那一半跟小的那一半
• 大的那半用Min Heap做，方便找最小值；小的那半用Max Heap做，方便找最大值
• 插入num時，一率先插入大的那半，再從大的那半找最小值塞回小的那半
• 判斷總數是偶數或奇數，決定中位數是單一數字或中間兩個平均

import heapq


class MedianFinder:
    def __init__(self):
        min_heap = []
        max_heap = []

        heapq.heapify(min_heap)
        heapq.heapify(max_heap)

        self.min_heap = min_heap
        self.max_heap = max_heap


    def addNum(self, num: int) -> None:
        heapq.heappush(self.min_heap, num)
        heapq.heappush(self.max_heap, -heapq.heappop(self.min_heap))

        if len(self.min_heap) < len(self.max_heap):
            heapq.heappush(self.min_heap, -heapq.heappop(self.max_heap))

    def findMedian(self) -> float:
        if len(self.min_heap) > len(self.max_heap):
            return self.min_heap[0]
        else:
            return (self.min_heap[0] + (-self.max_heap[0]) ) / 2