优先队列题目：数据流的中位数

题目

标题和出处

标题：数据流的中位数

出处：295. 数据流的中位数

难度

6 级

题目描述

要求

中位数是有序整数列表的中间的数。如果列表长度是偶数，则没有中间的数，此时中位数是中间的两个数的平均值。

• 例如，$\text{arr}=\text{[2, 3, 4]}$ 的中位数是 $\text{3}$。
• 例如，$\text{arr}=\text{[2, 3]}$ 的中位数是 $\frac{\text{3}}{\text{2}}=\text{2.5}$。

实现 $\text{MedianFinder}$ 类：

• $\text{MedianFinder()}$ 初始化 $\text{MedianFinder}$ 对象。
• $\text{void addNum(int num)}$ 将数据流中的整数 $\text{num}$ 添加到数据结构。
• $\text{double findMedian()}$ 返回目前所有元素的中位数。和正确答案的误差不超过 ${\text{10}}^{\text{-5}}$ 的答案被视为正确答案。

示例

示例 1：

输入：
$\text{["MedianFinder", "addNum", "addNum", "findMedian", "addNum", "findMedian"]}$
$\text{[[], [1], [2], [], [3], []]}$
输出：
$\text{[null, null, null, 1.5, null, 2.0]}$
解释：
$\text{MedianFinder medianFinder = new MedianFinder();}$
$\text{medianFinder.addNum(1);}$ // $\text{arr = [1]}$
$\text{medianFinder.addNum(2);}$ // $\text{arr = [1, 2]}$
$\text{medianFinder.findMedian();}$ // 返回 $\text{1.5}$
$\text{medianFinder.addNum(3);}$ // $\text{arr = [1, 2, 3]}$
$\text{medianFinder.findMedian();}$ // 返回 $\text{2.0}$

数据范围

• ${\text{-10}}^{\text{5}}\le \text{num}\le {\text{10}}^{\text{5}}$
• 在调用 $\text{findMedian}$ 方法之前，数据结构中至少有一个元素
• 最多调用 $\text{addNum}$ 和 $\text{findMedian}$ 方法 $\text{5}\times {\text{10}}^{\text{4}}$ 次

进阶

• 如果数据流中所有整数都在范围 $\text{[0, 100]}$ 内，你将如何优化你的算法？
• 如果数据流中 $\text{99\%}$ 的整数都在范围 $\text{[0, 100]}$ 内，你将如何优化你的算法？

解法

思路和算法

根据中位数的定义，中位数的值取决于有序列表的中间的一个数或两个数。为了得到数据流的中位数，需要设计一种满足以下要求的数据结构：

1. 该数据结构能维护元素的有序性（或者部分有序性），支持快速定位到中间元素；

2. 该数据结构支持快速添加元素，且添加元素之后仍能满足第 1 点要求。

满足上述要求的数据结构为优先队列，优先队列支持在对数时间复杂度内添加元素和在常数时间内获取队首元素。

根据优先队列的性质，为了得到数据流的中位数，需要创建两个优先队列分别存储较小的一半元素和较大的一半元素。存储较小的一半元素的优先队列是 $\text{lower}$，基于大顶堆实现，队首元素是优先队列中的最大元素；存储较大的一半元素的优先队列是 $\text{higher}$，基于小顶堆实现，队首元素是优先队列中的最小元素。两个优先队列满足以下要求：

• 优先队列 $\text{lower}$ 的队首元素小于等于优先队列 $\text{higher}$ 的队首元素；

• 优先队列 $\text{lower}$ 的元素个数大于等于优先队列 $\text{higher}$ 的元素个数；

• 两个优先队列的元素个数之差不超过 $1$。

为了满足上述要求，添加整数的操作如下：

• 如果数据流中的整数个数是偶数，首先将添加的整数加入 $\text{higher}$，然后将 $\text{higher}$ 的队首整数取出并将取出的整数加入 $\text{lower}$；

• 如果数据流中的整数个数是奇数，首先将添加的整数加入 $\text{lower}$，然后将 $\text{lower}$ 的队首整数取出并将取出的整数加入 $\text{higher}$。

当两个优先队列满足上述要求时，即可通过两个优先队列的队首元素计算中位数：

• 如果数据流中的整数个数是偶数，则两个优先队列的元素个数相等，中位数等于两个优先队列的队首元素的平均值；

• 如果数据流中的整数个数是奇数，则 $\text{lower}$ 比 $\text{higher}$ 多一个元素，中位数等于 $\text{lower}$ 的队首元素。

代码

class MedianFinder {
    PriorityQueue<Integer> lower;
    PriorityQueue<Integer> higher;
    int size;

    public MedianFinder() {
        lower = new PriorityQueue<Integer>(new Comparator<Integer>() {
            public int compare(Integer num1, Integer num2) {
                return num2 - num1;
            }
        });
        higher = new PriorityQueue<Integer>(new Comparator<Integer>() {
            public int compare(Integer num1, Integer num2) {
                return num1 - num2;
            }
        });
        size = 0;
    }

    public void addNum(int num) {
        if (size % 2 == 0) {
            higher.offer(num);
            lower.offer(higher.poll());
        } else {
            lower.offer(num);
            higher.offer(lower.poll());
        }
        size++;
    }

    public double findMedian() {
        return size % 2 == 0 ? (lower.peek() + higher.peek()) / 2.0 : lower.peek();
    }
}

复杂度分析

• 时间复杂度：构造方法和计算中位数操作的时间复杂度是 $O(1)$，添加整数操作的时间复杂度是 $O(\log n)$，其中 $n$ 是数据流中的整数个数。
  添加整数操作需要对其中一个优先队列进行添加元素和删除元素操作，对另一个优先队列进行添加元素操作，因此时间复杂度是 $O(\log n)$。
  计算中位数操作只要得到优先队列的队首元素即可计算，得到优先队列的队首元素的时间复杂度是 $O(1)$。

• 空间复杂度：$O(n)$，其中 $n$ 是数据流中的整数个数。空间复杂度主要取决于两个优先队列空间，两个优先队列存储数据流中的全部整数。

进阶问题答案

第一个进阶问题，如果数据流中所有整数都在范围 $[0,100]$ 内，则可以使用计数排序统计每个整数的个数，并根据数据流中的整数个数，使用双指针维护中位数。

第二个进阶问题，如果数据流中 $99\%$ 的整数都在范围 $[0,100]$ 内，则仍然可以使用计数排序统计每个整数的个数，并根据数据流中的整数个数，使用双指针维护中位数。对于小于 $0$ 或大于 $100$ 的整数，分别使用两个数组存储即可。如果中位数小于 $0$ 或大于 $100$，则只需要在对应的数组中计算中位数即可，由于大多数情况下中位数都在范围 $[0,100]$ 内，因此中位数小于 $0$ 或大于 $100$ 的情况是极少数，此时使用暴力解法不会影响总体时间复杂度。