算法的学习笔记—数据流中的中位数(牛客JZ41)

😀前言
在处理动态数据时，实时计算中位数是一个经典问题。中位数是排序后处于中间位置的数值，数据流中的中位数计算面临两个挑战：首先是数据量的动态变化，其次是需要保持元素的有序性。为了高效地解决这个问题，我们可以利用堆（Heap）结构。

🏠个人主页：尘觉主页

文章目录

😀数据流中的中位数

😀数据流中的中位数

题目链接

牛客网

😄问题描述

假设我们从数据流中不断读取数值，如果从中读取奇数个数值，中位数就是所有数值排序后位于中间的数值；如果读取偶数个数值，中位数就是排序后中间两个数值的平均值。

例如，输入数据流 [5, 2, 3, 4, 1, 6, 7, 0, 8]，我们需要依次计算数据流的中位数，得到的结果为：5.00 3.50 3.00 3.50 3.00 3.50 4.00 3.50 4.00。

💖示例1

输入：[5,2,3,4,1,6,7,0,8]

返回值："5.00 3.50 3.00 3.50 3.00 3.50 4.00 3.50 4.00 "

说明：数据流里面不断吐出的是5,2,3…,则得到的平均数分别为5,(5+2)/2,3…

💖示例2

输入：[1,1,1]

返回值："1.00 1.00 1.00 "

😊解题思路

为了高效地计算数据流的中位数，我们可以使用两个堆：大顶堆和小顶堆。

大顶堆（Max-Heap）：用于存储数据流中较小的一半元素，堆顶元素是这部分元素中的最大值。
小顶堆（Min-Heap）：用于存储数据流中较大的一半元素，堆顶元素是这部分元素中的最小值。

操作流程：

每当有新元素插入时，首先根据当前元素个数决定插入哪个堆。
- 如果总数为偶数，则新元素应插入小顶堆。但为了保证小顶堆的元素都大于大顶堆中的元素，我们可以先将新元素插入大顶堆，再将大顶堆堆顶的最大值弹出并插入小顶堆。
- 如果总数为奇数，则新元素应插入大顶堆。同样，为了保证大顶堆的元素都小于小顶堆的元素，我们可以先将新元素插入小顶堆，再将小顶堆堆顶的最小值弹出并插入大顶堆。
插入操作结束后，根据堆中元素的数量计算中位数：
- 如果总数为偶数，中位数为两个堆顶元素的平均值。
- 如果总数为奇数，中位数为小顶堆的堆顶元素。

🤔代码实现

以下是用 Java 实现的代码示例

import java.util.PriorityQueue;

public class MedianFinder {
    /* 
     * 大顶堆，存储数据流中较小的一半元素 
     * 比较器设为倒序，这样堆顶为最大值
     */
    private PriorityQueue<Integer> left = new PriorityQueue<>((o1, o2) -> o2 - o1);
    
    /* 
     * 小顶堆，存储数据流中较大的一半元素 
     * 默认是最小堆，堆顶为最小值
     */
    private PriorityQueue<Integer> right = new PriorityQueue<>();
    
    /* 当前数据流中元素的总数 */
    private int N = 0;

    /**
     * 插入新元素到数据流中
     * 
     * @param val 要插入的数据流中的新值
     */
    public void Insert(Integer val) {
        /* 插入操作：首先根据当前元素数量的奇偶性，决定插入哪个堆 */
        if (N % 2 == 0) {
            /*
             * 当 N 为偶数时：
             * - 先将新元素插入大顶堆（左半边），确保左半边堆顶最大。
             * - 然后将大顶堆中的最大值弹出并插入小顶堆（右半边），
             *   这样就保证了右半边的元素都大于左半边的元素。
             */
            left.add(val);
            right.add(left.poll());
        } else {
            /*
             * 当 N 为奇数时：
             * - 先将新元素插入小顶堆（右半边），确保右半边堆顶最小。
             * - 然后将小顶堆中的最小值弹出并插入大顶堆（左半边），
             *   这样就保证了左半边的元素都小于右半边的元素。
             */
            right.add(val);
            left.add(right.poll());
        }
        /* 插入完成后，元素总数加一 */
        N++;
    }

    /**
     * 获取当前数据流的中位数
     * 
     * @return 当前数据流的中位数
     */
    public Double GetMedian() {
        /* 
         * 判断元素总数的奇偶性来决定中位数的计算方式 
         * 如果总数为偶数，中位数为两个堆顶元素的平均值
         * 如果总数为奇数，中位数为小顶堆的堆顶元素
         */
        if (N % 2 == 0)
            return (left.peek() + right.peek()) / 2.0;
        else
            return (double) right.peek();
    }

    public static void main(String[] args) {
        MedianFinder mf = new MedianFinder();
        int[] nums = {5, 2, 3, 4, 1, 6, 7, 0, 8};
        
        /* 
         * 逐个插入元素到数据流，并输出当前的中位数 
         */
        for (int num : nums) {
            mf.Insert(num);
            System.out.print(mf.GetMedian() + " ");
        }
    }
}

代码解析

堆的构建：我们使用了 Java 的 PriorityQueue 类来实现大顶堆和小顶堆。大顶堆通过传入自定义比较器来实现，将较大元素优先弹出。小顶堆则是默认的最小堆。
插入逻辑：根据当前元素数量的奇偶性，决定元素插入的顺序。通过先插入一个堆，再将该堆顶元素移动到另一个堆，保证两个堆中的数据始终保持平衡。
获取中位数：根据元素的奇偶性，决定从哪个堆中获取中位数。如果元素总数为偶数，则取两个堆的堆顶元素平均值；如果为奇数，则直接返回小顶堆的堆顶元素。