Create 295-Find-Median-from-Data-Stream.py

Author: neetcode-gh

295-Find-Median-from-Data-Stream.py

@@ -0,0 +1,30 @@
+class MedianFinder:
+
+    def __init__(self):
+        """
+        initialize your data structure here.
+        """
+        # two heaps, large, small, minheap, maxheap
+        # heaps should be equal size
+        self.small, self.large = [], [] # maxHeap, minHeap (python default)
+
+    def addNum(self, num: int) -> None:
+        heapq.heappush(self.small, -1 * num)
+        
+        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 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.small) > len(self.large):
+            return -1 * self.small[0]
+        elif len(self.large) > len(self.small):
+            return self.large[0]
+        return (-1 * self.small[0] + self.large[0]) / 2