Hard

Find Median from Data Stream

091

Given a stream of numbers, return the median after each insertion.

EXAMPLES

Example 1

Input
{
  "nums": [
    1,
    2,
    3
  ]
}

Output
[
  1,
  1.5,
  2
]

FUNCTION SHAPE

nums: intArray→doubleArray

SOLUTION NOTE

Using a single SortedList completely trivializes this problem, but the two-heap solution is more commonly expected in interviews.

Reveal reference solution +

pythonREFERENCE

# Solution 1: Two Heaps
class MedianFinder:
    def __init__(self):
        self.left = []   # Max-heap (invert values)
        self.right = []  # Min-heap

    def addNum(self, num: int) -> None:
        heapq.heappush(self.left, -num)
        heapq.heappush(self.right, -heapq.heappop(self.left))
        if len(self.left) < len(self.right):
            heapq.heappush(self.left, -heapq.heappop(self.right))

    def findMedian(self) -> float:
        if len(self.left) > len(self.right):
            return -self.left[0]
        return (-self.left[0] + self.right[0]) / 2

# Solution 2: SortedList (trivializes the problem)
class MedianFinder:
    def __init__(self):
        self.arr = SortedList()

    def addNum(self, num: int) -> None:
        self.arr.add(num)

    def findMedian(self) -> float:
        n = len(self.arr)
        if n % 2 == 1:
            return self.arr[n // 2]
        return (self.arr[n // 2] + self.arr[n // 2 - 1]) / 2

TimeO(log n) per operation

SpaceO(n)

Open on LeetCode ↗