[Leetcode] 295. Find Median from Data Stream

https://leetcode.com/problems/find-median-from-data-stream/description/

initial

I thought it was just inputting in order to maxHeap -> minHeap and so on. But if we have a unsorted list, that doesnt work.

The most impt condition is that values stored in maxHeap are always smaller than that in minHeap.

import heapq

class MedianFinder:

    def __init__(self):
        self.maxHeap = []
        self.minHeap = []

    def addNum(self, num: int) -> None:
        if not self.maxHeap:
            heapq.heappush(self.maxHeap,-num)
        elif not self.minHeap:
            heapq.heappush(self.minHeap,num)

        if len(self.maxHeap)==len(self.minHeap):
            if num>self.minHeap[0]:
                val = heapq.heappop(self.minHeap)
                heapq.heappush(self.maxHeap,-val)
                heapq.heappush(self.minHeap, num)
    def findMedian(self) -> float:
        if not self.maxHeap:
            return None
        elif self.maxHeap and not self.minHeap:
            return - self.maxHeap[0]
            
        minVal = self.minHeap[0]
        maxVal = - self.maxHeap[0]
        if len(self.maxHeap)==len(self.minHeap):
            return (minVal+maxVal)/2
        else:
            return float(maxVal)

solution

So we have to ensure that
1) we always push to maxHeap first and take care of the difference in lengths of 2 heaps by popping and pushing
1.1) when the length difference exceeds by 2

like maxHeap = [0], minHeap=[] is fine but if maxHeap=[1,0],minHeap=[], that is when we need to take care of this.

Then we pop whichever is longer and push to the shorter one.

2) we have to consider unsorted list

import heapq

class MedianFinder:
    def __init__(self):
        # Max heap for the lower half (storing negative values)
        self.maxHeap = []
        # Min heap for the upper half
        self.minHeap = []

    def addNum(self, num: int) -> None:
        # Always push to max heap first
        heapq.heappush(self.maxHeap, -num)
        
        # Ensure that all elements in maxHeap are <= elements in minHeap
        if self.minHeap and -self.maxHeap[0] > self.minHeap[0]:
            # Move the top of maxHeap to minHeap
            val = -heapq.heappop(self.maxHeap)
            heapq.heappush(self.minHeap, val)
        
        # Rebalance heaps if their sizes differ by more than 1
        if len(self.maxHeap) > len(self.minHeap) + 1:
            val = -heapq.heappop(self.maxHeap)
            heapq.heappush(self.minHeap, val)
        
        if len(self.minHeap) > len(self.maxHeap) + 1:
            val = heapq.heappop(self.minHeap)
            heapq.heappush(self.maxHeap, -val)

    def findMedian(self) -> float:
        # If heaps are of equal length, return average of tops
        if len(self.maxHeap) == len(self.minHeap):
            return (-self.maxHeap[0] + self.minHeap[0]) / 2.0
        
        # Otherwise, return the top of the larger heap
        return -self.maxHeap[0] if len(self.maxHeap) > len(self.minHeap) else self.minHeap[0]

complexity

push and pop in heapq operations take log n time.
so addNum is log n time

n space cuz we are storing n ints.