Leetcode-and-Others/Leetcode 295 -- Find Median from Data Stream/Leetcode 295 -- Find Median from Data Stream.py at main · wena04/Leetcode-and-Others · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
# Leetcode 295: Find Median from Data Stream
# https://leetcode.com/problems/find-median-from-data-stream/

class MedianFinder:
    # use min and max heap to maintain inorder array

    def __init__(self):
        # small is maxheap, large is minheap since we want to get the median in O(1)
        # want to maintain heap so they are in +- 1 size of each other
        self.small, self.large = [], []

    def addNum(self, num: int) -> None:
        # python doesn't have maxheap so need to multiply -1 to make maxheap
        heapq.heappush(self.small, -1 * num)
        # if number in small is bigger than in large, move largest value to large
        if (self.small and self.large and (-1 * self.small[0]) > self.large[0]):
            val = -1 * heapq.heappop(self.small)
            heapq.heappush(self.large, val)

        # if the 2 array have uneven size, maintain it
        if len(self.small) > len(self.large) + 1:
            val = -1 * heapq.heappop(self.small)
            heapq.heappush(self.large, val)

        if len(self.large) > len(self.small) + 1:
            val = heapq.heappop(self.large)
            heapq.heappush(self.small, -1 * val)

    def findMedian(self) -> float:
        if len(self.large) > len(self.small):
            return self.large[0]
        if len(self.small) > len(self.large):
            return self.small[0] * -1
        return (self.large[0] + self.small[0] * -1) / 2

#time complexity: O(logN) for addNum, O(1) for findMedian
#space complexity: O(N)