leetcode hot100【LeetCode 295.数据流的中位数】java实现

LeetCode 295.数据流的中位数

题目描述

实现一个 MedianFinder 类：

• MedianFinder 类会维护一个数据流，并且提供以下两个方法：
  • addNum(num: int) -> None：将一个整数 num 添加到数据流中。
  • findMedian() -> float：返回目前所有元素的中位数。

中位数的定义：

• 如果数据的元素数量为奇数，则中位数是排序后正中间的元素。
• 如果数据的元素数量为偶数，则中位数是排序后两个中间元素的平均值。

示例 1：

输入：

["MedianFinder", "addNum", "addNum", "findMedian", "addNum", "findMedian"]
[[], [1], [2], [], [3], []]

输出：

[null, null, null, 1.5, null, 2.0]

解释：

• MedianFinder 初始化。
• addNum(1)：数据流是 [1]。
• addNum(2)：数据流是 [1, 2]。
• findMedian()：数据流是 [1, 2]，中位数是 (1 + 2) / 2 = 1.5。
• addNum(3)：数据流是 [1, 2, 3]。
• findMedian()：数据流是 [1, 2, 3]，中位数是 2。

Java 实现代码

import java.util.*;

public class MedianFinder {
    // 使用两个堆：一个大根堆（maxHeap）保存较小的一半，另一个小根堆（minHeap）保存较大的一半
    private PriorityQueue<Integer> maxHeap;
    private PriorityQueue<Integer> minHeap;

    /** 初始化数据流 */
    public MedianFinder() {
        // 大根堆：存储较小的一半
        maxHeap = new PriorityQueue<>((a, b) -> b - a);
        // 小根堆：存储较大的一半
        minHeap = new PriorityQueue<>();
    }
    
    /** 添加数字到数据流 */
    public void addNum(int num) {
        // 确保maxHeap的最大元素小于minHeap的最小元素
        if (maxHeap.isEmpty() || num <= maxHeap.peek()) {
            maxHeap.offer(num);
        } else {
            minHeap.offer(num);
        }

        // 调整堆的大小，确保两个堆的大小差距不超过1
        if (maxHeap.size() > minHeap.size() + 1) {
            minHeap.offer(maxHeap.poll());
        } else if (minHeap.size() > maxHeap.size()) {
            maxHeap.offer(minHeap.poll());
        }
    }
    
    /** 查找数据流的中位数 */
    public double findMedian() {
        // 如果两个堆的大小相等，返回两个堆顶元素的平均值
        if (maxHeap.size() == minHeap.size()) {
            return (maxHeap.peek() + minHeap.peek()) / 2.0;
        } else {
            // 如果maxHeap更多，返回maxHeap的堆顶元素
            return maxHeap.peek();
        }
    }
}

解题思路

1. 使用两个堆 为了高效找到中位数，我们可以使用两个堆来维护数据流的一半：

   • 大根堆（maxHeap）：存储数据流中较小的一半元素。
   • 小根堆（minHeap）：存储数据流中较大的一半元素。
     这样可以保证：
   • maxHeap 中的最大元素小于 minHeap 中的最小元素。
   • 通过维护堆的大小差异，maxHeap 和 minHeap 的元素数量要么相等，要么 maxHeap 至少比 minHeap 多一个元素。

2. 插入新元素

   • 每次插入一个新元素时，我们首先判断该元素应该放入哪个堆。如果它小于等于 maxHeap 的堆顶元素，则插入 maxHeap，否则插入 minHeap。
   • 然后，我们根据堆的大小调整两堆的元素数量，确保两堆的大小差不超过1。

3. 计算中位数

   • 如果两个堆的大小相等，说明数据流的元素个数为偶数，返回两个堆顶元素的平均值。
   • 如果 maxHeap 的大小较大，则中位数就是 maxHeap 的堆顶元素。

复杂度分析

• 时间复杂度: addNum: O(logn)，其中 n 为累计添加的数的数量。findMedian: O(1)。
• 空间复杂度: O(n)，主要为优先队列的开销。

举例说明执行过程

假设输入为：[1, 2, 3]，并且我们依次调用 addNum(num) 和 findMedian()。

1. 初始化：

   • maxHeap = []
   • minHeap = []

2. 调用 addNum(1)：

   • 1 被插入 maxHeap，因为 maxHeap 为空。
   • maxHeap = [1]
   • minHeap = []

3. 调用 addNum(2)：

   • 2 被插入 minHeap，因为 2 > 1（maxHeap 的堆顶）。
   • maxHeap = [1]
   • minHeap = [2]

4. 调用 findMedian()：

   • maxHeap.size() == minHeap.size()，返回 (1 + 2) / 2 = 1.5。

5. 调用 addNum(3)：

   • 3 被插入 minHeap，因为 3 > 2（minHeap 的堆顶）。
   • maxHeap = [1]
   • minHeap = [2, 3]
   • 调整堆，maxHeap 的堆顶元素移至 minHeap，最终 maxHeap = [2] 和 minHeap = [3]。

6. 调用 findMedian()：

   • maxHeap.size() > minHeap.size()，返回 maxHeap 的堆顶元素 2。

最终输出结果是：[null, null, null, 1.5, null, 2.0]。