Hard

题目描述

中位数是有序整数列表中的中间值。如果列表的大小是偶数,则没有中间值,中位数是两个中间值的平均值。

  • 例如,对于 arr = [2,3,4],中位数是 3。
  • 例如,对于 arr = [2,3],中位数是 (2 + 3) / 2 = 2.5。

实现 MedianFinder 类:

  • MedianFinder() 初始化 MedianFinder 对象。
  • void addNum(int num) 从数据流中添加一个整数 num 到数据结构中。
  • double findMedian() 返回目前所有元素的中位数。与实际答案误差在 10^-5 内的答案都可以接受。

示例 1:

输入
["MedianFinder", "addNum", "addNum", "findMedian", "addNum", "findMedian"]
[[], [1], [2], [], [3], []]
输出
[null, null, null, 1.5, null, 2.0]

解释
MedianFinder medianFinder = new MedianFinder();
medianFinder.addNum(1);    // arr = [1]
medianFinder.addNum(2);    // arr = [1, 2]
medianFinder.findMedian(); // return 1.5 (i.e., (1 + 2) / 2)
medianFinder.addNum(3);    // arr = [1, 2, 3]
medianFinder.findMedian(); // return 2.0

约束条件:

  • -10^5 <= num <= 10^5
  • 在调用 findMedian 之前,数据结构中至少有一个元素。
  • 最多会对 addNum 和 findMedian 进行 5 * 10^4 次调用。

进阶:

  • 如果数据流中所有整数都在 [0, 100] 范围内,你将如何优化你的算法?
  • 如果数据流中 99% 的整数都在 [0, 100] 范围内,你将如何优化你的算法?

解题思路

解题思路

这道题要求在数据流中动态维护中位数,有几种常见的解法:

  1. 朴素解法:使用数组存储所有数据,每次查询中位数时排序。时间复杂度 O(n log n),不够高效。

  2. 有序数组 + 二分插入:维护一个有序数组,插入时使用二分查找定位位置。插入时间复杂度 O(n),查询 O(1)。

  3. 双堆解法(推荐):使用两个堆来维护数据的左半部分和右半部分。

    • 左半部分用大顶堆(max heap)存储较小的一半数据
    • 右半部分用小顶堆(min heap)存储较大的一半数据
    • 保持两个堆的大小平衡(差值不超过1)

双堆解法的核心思想是:

  • 如果总数是奇数,让大顶堆多存一个元素,中位数就是大顶堆的堆顶
  • 如果总数是偶数,中位数就是两个堆顶的平均值

添加数字的策略:

  1. 如果大顶堆为空或数字 <= 大顶堆堆顶,加入大顶堆
  2. 否则加入小顶堆
  3. 调整两个堆的大小使其平衡

这种方法的时间复杂度:添加 O(log n),查询 O(1),空间复杂度 O(n)。

代码实现

class MedianFinder {
private:
    priority_queue<int> maxHeap; // 存储较小的一半
    priority_queue<int, vector<int>, greater<int>> minHeap; // 存储较大的一半
    
public:
    MedianFinder() {
        
    }
    
    void addNum(int num) {
        // 如果大顶堆为空或者num小于等于大顶堆的堆顶,添加到大顶堆
        if (maxHeap.empty() || num <= maxHeap.top()) {
            maxHeap.push(num);
        } else {
            minHeap.push(num);
        }
        
        // 平衡两个堆的大小
        if (maxHeap.size() > minHeap.size() + 1) {
            minHeap.push(maxHeap.top());
            maxHeap.pop();
        } else if (minHeap.size() > maxHeap.size() + 1) {
            maxHeap.push(minHeap.top());
            minHeap.pop();
        }
    }
    
    double findMedian() {
        if (maxHeap.size() > minHeap.size()) {
            return maxHeap.top();
        } else if (minHeap.size() > maxHeap.size()) {
            return minHeap.top();
        } else {
            return (maxHeap.top() + minHeap.top()) / 2.0;
        }
    }
};
import heapq

class MedianFinder:

    def __init__(self):
        self.max_heap = []  # 存储较小的一半(使用负数模拟大顶堆)
        self.min_heap = []  # 存储较大的一半(小顶堆)

    def addNum(self, num: int) -> None:
        # 如果大顶堆为空或者num小于等于大顶堆的堆顶,添加到大顶堆
        if not self.max_heap or num <= -self.max_heap[0]:
            heapq.heappush(self.max_heap, -num)
        else:
            heapq.heappush(self.min_heap, num)
        
        # 平衡两个堆的大小
        if len(self.max_heap) > len(self.min_heap) + 1:
            val = -heapq.heappop(self.max_heap)
            heapq.heappush(self.min_heap, val)
        elif len(self.min_heap) > len(self.max_heap) + 1:
            val = heapq.heappop(self.min_heap)
            heapq.heappush(self.max_heap, -val)

    def findMedian(self) -> float:
        if len(self.max_heap) > len(self.min_heap):
            return -self.max_heap[0]
        elif len(self.min_heap) > len(self.max_heap):
            return self.min_heap[0]
        else:
            return (-self.max_heap[0] + self.min_heap[0]) / 2.0
public class MedianFinder {
    private PriorityQueue<int, int> maxHeap; // 存储较小的一半
    private PriorityQueue<int, int> minHeap; // 存储较大的一半

    public MedianFinder() {
        maxHeap = new PriorityQueue<int, int>(Comparer<int>.Create((a, b) => b.CompareTo(a))); // 大顶堆
        minHeap = new PriorityQueue<int, int>(); // 小顶堆
    }
    
    public void AddNum(int num) {
        // 如果大顶堆为空或者num小于等于大顶堆的堆顶,添加到大顶堆
        if (maxHeap.Count == 0 || num <= maxHeap.Peek()) {
            maxHeap.Enqueue(num, num);
        } else {
            minHeap.Enqueue(num, num);
        }
        
        // 平衡两个堆的大小
        if (maxHeap.Count > minHeap.Count + 1) {
            int val = maxHeap.Dequeue();
            minHeap.Enqueue(val, val);
        } else if (minHeap.Count > maxHeap.Count + 1) {
            int val = minHeap.Dequeue();
            maxHeap.Enqueue(val, val);
        }
    }
    
    public double FindMedian() {
        if (maxHeap.Count > minHeap.Count) {
            return maxHeap.Peek();
        } else if (minHeap.Count > maxHeap.Count) {
            return minHeap.Peek();
        } else {
            return (maxHeap.Peek() + minHeap.Peek()) / 2.0;
        }
    }
}
var MedianFinder = function() {
    this.maxHeap = new MaxPriorityQueue(); // 存储较小的一半
    this.minHeap = new MinPriorityQueue(); // 存储较大的一半
};

MedianFinder.prototype.addNum = function(num) {
    // 如果大顶堆为空或者num小于等于大顶堆的堆顶,添加到大顶堆
    if (this.maxHeap.isEmpty() || num <= this.maxHeap.front().element) {
        this.maxHeap.enqueue(num);
    } else {
        this.minHeap.enqueue(num);
    }
    
    // 平衡两个堆的大小
    if (this.maxHeap.size() > this.minHeap.size() + 1) {
        let val = this.maxHeap.dequeue().element;
        this.minHeap.enqueue(val);
    } else if (this.minHeap.size() > this.maxHeap.size() + 1) {
        let val = this.minHeap.dequeue().element;
        this.maxHeap.enqueue(val);
    }
};

MedianFinder.prototype.findMedian = function() {
    if (this.maxHeap.size() > this.minHeap.size()) {
        return this.maxHeap.front().element;
    } else if (this.minHeap.size() > this.maxHeap.size()) {
        return this.minHeap.front().element;
    } else {
        return (this.maxHeap.front().element + this.minHeap.front().element) / 2.0;
    }
};

复杂度分析

解法时间复杂度空间复杂度
双堆解法addNum: O(log n), findMedian: O(1)O(n)
有序数组addNum: O(n), findMedian: O(1)O(n)
朴素排序addNum: O(1), findMedian: O(n log n)O(n)

推荐解法:双堆解法,在频繁插入和查询的场景下表现最优。

相关题目