295. Find Median from Data Stream

We can find the median from a data stream by always keeping a sorted list of all the items so far. We then can always find the median quickly ($O(\log n)$ time).

The two heaps approach is better though time complexity wise. It’s $O(1)$ to find medians, unlike this sorted list which is $O(\log n)$. Insertion is the same for both though.

from sortedcontainers import SortedList
 
class MedianFinder:
    def __init__(self):
        self.nums = SortedList()
        self.len = 0
 
    def addNum(self, num: int) -> None:
        self.nums.add(num)
        self.len += 1
 
    def findMedian(self) -> float:
        if self.len % 2:
            return self.nums[self.len // 2]
        else:
            left = self.nums[self.len // 2 - 1]
            right = self.nums[self.len // 2]
            return (left + right) / 2

import heapq
 
class MedianFinder:
    def __init__(self):
        self.small = []  # max-heap (store negative numbers)
        self.large = []  # min-heap
 
    def addNum(self, num: int) -> None:
        # Push to small (max-heap)
        heapq.heappush(self.small, -num)
 
        # Ensure every element in small <= every element in large
        if self.small and self.large and -self.small[0] > self.large[0]:
            val = -heapq.heappop(self.small)
            heapq.heappush(self.large, val)
 
        # Balance sizes (small can have at most 1 more element than large)
        if len(self.small) > len(self.large) + 1:
            heapq.heappush(self.large, -heapq.heappop(self.small))
        if len(self.large) > len(self.small):
            heapq.heappush(self.small, -heapq.heappop(self.large))
 
    def findMedian(self) -> float:
        if len(self.small) > len(self.large):
            return -self.small[0]
        return (-self.small[0] + self.large[0]) / 2