Find Median from Data Stream LeetCode solution in Python

1
class MedianFinder:
2

3
    def __init__(self):
4
        ### max heap to store the first half of the list
5
        self.maxHeap = []
6
        ### min heap to store the second half of the list
7
        self.minHeap = []
8

9
    def addNum(self, num: int) -> None:
10
        ### push num into the correct heap
11
        if not self.maxHeap or num <= -self.maxHeap[0]:
12
            heappush(self.maxHeap, -num)
13
        else:
14
            heappush(self.minHeap, num)
15

16
        ### banance the two heaps so that each of them representing half of the list
17
        ### for odd length list, len(maxHeap) == len(minHeap)+1
18
        ### for even length list, len(maxHeap) == len(minHeap)
19
        if len(self.minHeap) > len(self.maxHeap):
20
            heappush(self.maxHeap, -heappop(self.minHeap))
21
        elif len(self.maxHeap) > len(self.minHeap) + 1:
22
            heappush(self.minHeap, -heappop(self.maxHeap))
23

24
    def findMedian(self) -> float:
25

26
        ### if the length of entire list is even,
27
        ### get the mean of the two middle values
28
        if (len(self.maxHeap) + len(self.minHeap)) % 2 == 0:
29
            return (-self.maxHeap[0] + self.minHeap[0]) / 2
30

31
        ### when odd, we know that the median is in maxHeap
32
        return -self.maxHeap[0]