Skip to content

295. Find Median from Data Stream 👍

  • Time: $O(n\log n)$
  • Space: $O(n)$
 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
class MedianFinder {
 public:
  void addNum(int num) {
    if (maxHeap.empty() || num <= maxHeap.top())
      maxHeap.push(num);
    else
      minHeap.push(num);

    // Balance the two heaps s.t.
    // |maxHeap| >= |minHeap| and |maxHeap| - |minHeap| <= 1.
    if (maxHeap.size() < minHeap.size())
      maxHeap.push(minHeap.top()), minHeap.pop();
    else if (maxHeap.size() - minHeap.size() > 1)
      minHeap.push(maxHeap.top()), maxHeap.pop();
  }

  double findMedian() {
    if (maxHeap.size() == minHeap.size())
      return (maxHeap.top() + minHeap.top()) / 2.0;
    return maxHeap.top();
  }

 private:
  priority_queue<int> maxHeap;
  priority_queue<int, vector<int>, greater<>> minHeap;
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class MedianFinder {
  public void addNum(int num) {
    if (maxHeap.isEmpty() || num <= maxHeap.peek())
      maxHeap.offer(num);
    else
      minHeap.offer(num);

    // Balance the two heaps s.t.
    // |maxHeap| >= |minHeap| and |maxHeap| - |minHeap| <= 1.
    if (maxHeap.size() < minHeap.size())
      maxHeap.offer(minHeap.poll());
    else if (maxHeap.size() - minHeap.size() > 1)
      minHeap.offer(maxHeap.poll());
  }

  public double findMedian() {
    if (maxHeap.size() == minHeap.size())
      return (double) (maxHeap.peek() + minHeap.peek()) / 2.0;
    return (double) maxHeap.peek();
  }

  private Queue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder());
  private Queue<Integer> minHeap = new PriorityQueue<>();
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
class MedianFinder:
  def __init__(self):
    self.maxHeap = []
    self.minHeap = []

  def addNum(self, num: int) -> None:
    if not self.maxHeap or num <= -self.maxHeap[0]:
      heapq.heappush(self.maxHeap, -num)
    else:
      heapq.heappush(self.minHeap, num)

    # Balance the two heaps s.t.
    # |maxHeap| >= |minHeap| and |maxHeap| - |minHeap| <= 1.
    if len(self.maxHeap) < len(self.minHeap):
      heapq.heappush(self.maxHeap, -heapq.heappop(self.minHeap))
    elif len(self.maxHeap) - len(self.minHeap) > 1:
      heapq.heappush(self.minHeap, -heapq.heappop(self.maxHeap))

  def findMedian(self) -> float:
    if len(self.maxHeap) == len(self.minHeap):
      return (-self.maxHeap[0] + self.minHeap[0]) / 2.0
    return -self.maxHeap[0]