41 数据流中的中位数

Lin2J
2021-07-21 / 0 评论 / 86 阅读
温馨提示:
本文最后更新于2021-07-21,若内容或图片失效,请留言反馈。

剑指 Offer 41. 数据流中的中位数

如何得到一个数据流中的中位数?如果从数据流中读出奇数个数值,那么中位数就是所有数值排序之后位于中间的数值。如果从数据流中读出偶数个数值,那么中位数就是所有数值排序之后中间两个数的平均值。

解题思路:

中位数又叫中值,是按顺序排列的一组数据中居于中间位置的数。

因为需要顺序排列,所以我们需要对数据流中的数字进行排序。

自定义规则

以中位数为界限,我们将小于中位数的数字成为较小数,大于中位数的数字成为较大数。较小数的那部分称为较小部分,较大数的部分称为较大部分。然后规定添加过程中,要保持较大部分的个数大于等于较小部分的个数。

我们在计算中位数的时候,需要分为两种情况,元素总个数为奇数和偶数:

  • 元素总数为奇数时,需要拿到较大部分中的最小数作为中位数返回;
  • 元素总数为偶数时,需要拿到较大部分的最小数和叫小部分的最大数的平均数返回;

数据结构选型

接下来就是思考如何存储较小部分,然后可以快速拿到较小部分的最大数,以及如何存储较大部分,然后可以快速拿到较大部分的最小数?

  • 可以使用大根堆来存储较小部分,因为大根堆的根是所有堆内元素最大的数字;
  • 可以使用小根堆来存储较大部分,因为小根堆的根是所有堆内元素最小的数字;

思路梳理

  1. 将数据流的数据,分为较大部分(用小根堆 $A$ 存储)和较小部分(用大根堆 $B$ 存储)。
  2. 添加数字 $x$ 的时候
    1. 如果此时总数为偶数,那么根据前面的规定,添加后 $A$ 要比 $B$ 多一个数字,先将 $x$ 添加到 $B$ 中,然后将 $B$ 的根添加到 $A$ 中。
    2. 如果此时总数为奇数,此时 $A.size > B.size$,那么为了维持 $A$ 和 $B$ 的大小平衡,添加后 $A$ 和 $B$ 的个数要想等,先将 $x$ 添加到 $A$ 中,然后将 $A$ 的根添加到 $B$ 中。
    3. 总数加一。
  3. 求中位数时
    1. 如果总数为奇数,那么直接返回 $A$ 的根。
    2. 如果总数为偶数,那么返回 $A$ 和 $B$ 的根的平均数。

代码:

class MedianFinder {
    /**
     * 小顶堆存放的是数据流排序后的后半部分,即较大的一半
     */
    private final PriorityQueue<Integer> minHeap;
    /**
     * 大顶堆存放的是数据流排序后的前半部分,即较小的一半
     */
    private final PriorityQueue<Integer> maxHeap;

    private int total = 0;

    /**
     * initialize your data structure here.
     */
    public MedianFinder() {
        minHeap = new PriorityQueue<>(Comparator.naturalOrder());
        maxHeap = new PriorityQueue<>(Comparator.reverseOrder());
    }

    /**
     * 添加的过程中,要保持小顶堆的元素个数大于等于大顶堆的个数,
     * 即总元素个数为奇数的时候,小顶堆要比大顶堆多一个元素
     */
    public void addNum(int num) {
        // 如果是偶数,则入小顶堆,但是不能直接入
        // 原因是,这个数字可能属于较小的一半,所以需要先入大顶堆,再把大顶堆的根弹出
        if ((total & 1) == 0) {
            maxHeap.offer(num);
            minHeap.offer(maxHeap.poll());
        } else {
            // 如果是奇数,则入大顶堆,但是不能直接入
            // 原因是,这个数字可能属于较大的一半,所以需要先入小顶堆,再把小顶堆的根弹出
            minHeap.offer(num);
            maxHeap.offer(minHeap.poll());
        }
        total++;
    }

    public double findMedian() {
        // 偶数则返回中间两个元素的和的一半
        if ((total & 1) == 0) {
            return (double) (minHeap.peek() + maxHeap.peek()) / 2;
        } else {
            // 奇数则返回小顶堆的根
            return (double) minHeap.peek();
        }
    }
}

复杂度分析

  • 时间复杂度

    • 添加数字 $O(logN)$:堆的弹出和压入使用 $O(logN)$ 的时间;
    • 求中位数$O(1)$:获取堆顶元素使用 $O(1)$ 时间。
  • 空间复杂度$O(N)$:其中 $N$ 为数据流中的元素数量,小顶堆 $A$ 和大顶堆 $B$ 最多同时保存 $N$ 个元素。