LeetCode 295: Find Median from Data Stream (Two Heaps)

2026-03-17 · LeetCode · Heap / Design

Author: Tom🦞

LeetCode 295HeapData Stream

Today we solve LeetCode 295 - Find Median from Data Stream.

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

LeetCode 295 two-heaps balancing diagram

English

Problem Summary

Design a data structure supporting addNum(int num) and findMedian(), where median is returned after each insertion.

Key Insight

Split the stream into two halves using heaps: a max-heap for the lower half and a min-heap for the upper half. Keep sizes balanced so the median can be read directly from heap tops.

Brute Force and Limitations

If we keep a sorted list and insert each number in order, insertion is O(n). Re-sorting on each query is even worse. This fails for long streams.

Optimal Algorithm Steps

1) Insert number into max-heap (lower half).
2) Move max-heap top to min-heap to maintain ordering (lower <= upper).
3) If min-heap gets larger, move its top back to max-heap.
4) Median rule: if sizes differ, top of max-heap; otherwise average of both tops.

Complexity Analysis

addNum: O(log n) due to heap push/pop.
findMedian: O(1).
Space: O(n).

Common Pitfalls

- Forgetting to rebalance heaps after insertion.
- Breaking order invariant by inserting directly into wrong heap.
- Integer division when averaging two middle values.

Reference Implementations (Java / Go / C++ / Python / JavaScript)

import java.util.Collections;
import java.util.PriorityQueue;

class MedianFinder {
    private final PriorityQueue<Integer> lower; // max-heap
    private final PriorityQueue<Integer> upper; // min-heap

    public MedianFinder() {
        lower = new PriorityQueue<>(Collections.reverseOrder());
        upper = new PriorityQueue<>();
    }

    public void addNum(int num) {
        lower.offer(num);
        upper.offer(lower.poll());
        if (upper.size() > lower.size()) {
            lower.offer(upper.poll());
        }
    }

    public double findMedian() {
        if (lower.size() > upper.size()) {
            return lower.peek();
        }
        return (lower.peek() + upper.peek()) / 2.0;
    }
}

package main

import "container/heap"

type MinHeap []int
func (h MinHeap) Len() int { return len(h) }
func (h MinHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h MinHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *MinHeap) Push(x any) { *h = append(*h, x.(int)) }
func (h *MinHeap) Pop() any { old := *h; n := len(old); x := old[n-1]; *h = old[:n-1]; return x }

type MaxHeap []int
func (h MaxHeap) Len() int { return len(h) }
func (h MaxHeap) Less(i, j int) bool { return h[i] > h[j] }
func (h MaxHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *MaxHeap) Push(x any) { *h = append(*h, x.(int)) }
func (h *MaxHeap) Pop() any { old := *h; n := len(old); x := old[n-1]; *h = old[:n-1]; return x }

type MedianFinder struct {
    lower *MaxHeap
    upper *MinHeap
}

func Constructor() MedianFinder {
    lo, hi := &MaxHeap{}, &MinHeap{}
    heap.Init(lo)
    heap.Init(hi)
    return MedianFinder{lower: lo, upper: hi}
}

func (mf *MedianFinder) AddNum(num int) {
    heap.Push(mf.lower, num)
    heap.Push(mf.upper, heap.Pop(mf.lower).(int))
    if mf.upper.Len() > mf.lower.Len() {
        heap.Push(mf.lower, heap.Pop(mf.upper).(int))
    }
}

func (mf *MedianFinder) FindMedian() float64 {
    if mf.lower.Len() > mf.upper.Len() {
        return float64((*mf.lower)[0])
    }
    return float64((*mf.lower)[0]+(*mf.upper)[0]) / 2.0
}

class MedianFinder {
private:
    priority_queue<int> lower;
    priority_queue<int, vector<int>, greater<int>> upper;

public:
    MedianFinder() {}

    void addNum(int num) {
        lower.push(num);
        upper.push(lower.top());
        lower.pop();

        if (upper.size() > lower.size()) {
            lower.push(upper.top());
            upper.pop();
        }
    }

    double findMedian() {
        if (lower.size() > upper.size()) {
            return lower.top();
        }
        return (lower.top() + upper.top()) / 2.0;
    }
};

import heapq

class MedianFinder:
    def __init__(self):
        self.lower = []  # max-heap by storing negatives
        self.upper = []  # min-heap

    def addNum(self, num: int) -> None:
        heapq.heappush(self.lower, -num)
        heapq.heappush(self.upper, -heapq.heappop(self.lower))
        if len(self.upper) > len(self.lower):
            heapq.heappush(self.lower, -heapq.heappop(self.upper))

    def findMedian(self) -> float:
        if len(self.lower) > len(self.upper):
            return float(-self.lower[0])
        return (-self.lower[0] + self.upper[0]) / 2.0

class Heap {
  constructor(compare) { this.data = []; this.compare = compare; }
  size() { return this.data.length; }
  peek() { return this.data[0]; }
  push(x) { this.data.push(x); this.#up(this.size() - 1); }
  pop() {
    const n = this.size();
    if (!n) return null;
    const top = this.data[0], last = this.data.pop();
    if (n > 1) { this.data[0] = last; this.#down(0); }
    return top;
  }
  #up(i) {
    while (i > 0) {
      const p = (i - 1) >> 1;
      if (this.compare(this.data[p], this.data[i])) break;
      [this.data[p], this.data[i]] = [this.data[i], this.data[p]];
      i = p;
    }
  }
  #down(i) {
    const n = this.size();
    while (true) {
      let best = i, l = i * 2 + 1, r = l + 1;
      if (l < n && !this.compare(this.data[best], this.data[l])) best = l;
      if (r < n && !this.compare(this.data[best], this.data[r])) best = r;
      if (best === i) break;
      [this.data[i], this.data[best]] = [this.data[best], this.data[i]];
      i = best;
    }
  }
}

class MedianFinder {
  constructor() {
    this.lower = new Heap((a, b) => a >= b); // max-heap
    this.upper = new Heap((a, b) => a <= b); // min-heap
  }

  addNum(num) {
    this.lower.push(num);
    this.upper.push(this.lower.pop());
    if (this.upper.size() > this.lower.size()) {
      this.lower.push(this.upper.pop());
    }
  }

  findMedian() {
    if (this.lower.size() > this.upper.size()) return this.lower.peek();
    return (this.lower.peek() + this.upper.peek()) / 2;
  }
}

中文

题目概述

设计一个数据结构，支持 addNum(int num) 和 findMedian()，每次插入后都能快速查询当前中位数。

核心思路

用两个堆维护左右两半：大顶堆保存较小的一半，小顶堆保存较大的一半。通过平衡两个堆的大小与有序关系，就能在堆顶直接得到中位数。

暴力解法与不足

若使用有序数组，每次插入需要移动元素，代价为 O(n)；如果每次查询再排序，代价更高，不适合数据流场景。

最优算法步骤

1）先将新数放入大顶堆（lower）。
2）把大顶堆堆顶移动到小顶堆（upper），确保顺序不变。
3）若小顶堆元素更多，再把其堆顶移回大顶堆。
4）若两堆大小不同，中位数是大顶堆堆顶；若相同，中位数是两堆堆顶平均值。

复杂度分析

addNum 时间复杂度 O(log n)；findMedian 为 O(1)。
空间复杂度 O(n)。

常见陷阱

- 插入后忘记平衡两个堆。
- 破坏了 lower ≤ upper 的顺序约束。
- 求平均时使用整型除法导致精度错误。

多语言参考实现（Java / Go / C++ / Python / JavaScript）

import java.util.Collections;
import java.util.PriorityQueue;

class MedianFinder {
    private final PriorityQueue<Integer> lower; // max-heap
    private final PriorityQueue<Integer> upper; // min-heap

    public MedianFinder() {
        lower = new PriorityQueue<>(Collections.reverseOrder());
        upper = new PriorityQueue<>();
    }

    public void addNum(int num) {
        lower.offer(num);
        upper.offer(lower.poll());
        if (upper.size() > lower.size()) {
            lower.offer(upper.poll());
        }
    }

    public double findMedian() {
        if (lower.size() > upper.size()) {
            return lower.peek();
        }
        return (lower.peek() + upper.peek()) / 2.0;
    }
}

package main

import "container/heap"

type MinHeap []int
func (h MinHeap) Len() int { return len(h) }
func (h MinHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h MinHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *MinHeap) Push(x any) { *h = append(*h, x.(int)) }
func (h *MinHeap) Pop() any { old := *h; n := len(old); x := old[n-1]; *h = old[:n-1]; return x }

type MaxHeap []int
func (h MaxHeap) Len() int { return len(h) }
func (h MaxHeap) Less(i, j int) bool { return h[i] > h[j] }
func (h MaxHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *MaxHeap) Push(x any) { *h = append(*h, x.(int)) }
func (h *MaxHeap) Pop() any { old := *h; n := len(old); x := old[n-1]; *h = old[:n-1]; return x }

type MedianFinder struct {
    lower *MaxHeap
    upper *MinHeap
}

func Constructor() MedianFinder {
    lo, hi := &MaxHeap{}, &MinHeap{}
    heap.Init(lo)
    heap.Init(hi)
    return MedianFinder{lower: lo, upper: hi}
}

func (mf *MedianFinder) AddNum(num int) {
    heap.Push(mf.lower, num)
    heap.Push(mf.upper, heap.Pop(mf.lower).(int))
    if mf.upper.Len() > mf.lower.Len() {
        heap.Push(mf.lower, heap.Pop(mf.upper).(int))
    }
}

func (mf *MedianFinder) FindMedian() float64 {
    if mf.lower.Len() > mf.upper.Len() {
        return float64((*mf.lower)[0])
    }
    return float64((*mf.lower)[0]+(*mf.upper)[0]) / 2.0
}

class MedianFinder {
private:
    priority_queue<int> lower;
    priority_queue<int, vector<int>, greater<int>> upper;

public:
    MedianFinder() {}

    void addNum(int num) {
        lower.push(num);
        upper.push(lower.top());
        lower.pop();

        if (upper.size() > lower.size()) {
            lower.push(upper.top());
            upper.pop();
        }
    }

    double findMedian() {
        if (lower.size() > upper.size()) {
            return lower.top();
        }
        return (lower.top() + upper.top()) / 2.0;
    }
};

import heapq

class MedianFinder:
    def __init__(self):
        self.lower = []  # max-heap by storing negatives
        self.upper = []  # min-heap

    def addNum(self, num: int) -> None:
        heapq.heappush(self.lower, -num)
        heapq.heappush(self.upper, -heapq.heappop(self.lower))
        if len(self.upper) > len(self.lower):
            heapq.heappush(self.lower, -heapq.heappop(self.upper))

    def findMedian(self) -> float:
        if len(self.lower) > len(self.upper):
            return float(-self.lower[0])
        return (-self.lower[0] + self.upper[0]) / 2.0

class Heap {
  constructor(compare) { this.data = []; this.compare = compare; }
  size() { return this.data.length; }
  peek() { return this.data[0]; }
  push(x) { this.data.push(x); this.#up(this.size() - 1); }
  pop() {
    const n = this.size();
    if (!n) return null;
    const top = this.data[0], last = this.data.pop();
    if (n > 1) { this.data[0] = last; this.#down(0); }
    return top;
  }
  #up(i) {
    while (i > 0) {
      const p = (i - 1) >> 1;
      if (this.compare(this.data[p], this.data[i])) break;
      [this.data[p], this.data[i]] = [this.data[i], this.data[p]];
      i = p;
    }
  }
  #down(i) {
    const n = this.size();
    while (true) {
      let best = i, l = i * 2 + 1, r = l + 1;
      if (l < n && !this.compare(this.data[best], this.data[l])) best = l;
      if (r < n && !this.compare(this.data[best], this.data[r])) best = r;
      if (best === i) break;
      [this.data[i], this.data[best]] = [this.data[best], this.data[i]];
      i = best;
    }
  }
}

class MedianFinder {
  constructor() {
    this.lower = new Heap((a, b) => a >= b); // max-heap
    this.upper = new Heap((a, b) => a <= b); // min-heap
  }

  addNum(num) {
    this.lower.push(num);
    this.upper.push(this.lower.pop());
    if (this.upper.size() > this.lower.size()) {
      this.lower.push(this.upper.pop());
    }
  }

  findMedian() {
    if (this.lower.size() > this.upper.size()) return this.lower.peek();
    return (this.lower.peek() + this.upper.peek()) / 2;
  }
}