数据流的中位数 [双堆&桶的使用]

482 阅读2分钟

我报名参加金石计划1期挑战——瓜分10万奖池,这是我的第3篇文章,数据流的中位数 [双堆&桶的使用] - 掘金 (juejin.cn)

前言

在不断添加元素的情况下,获取数据中的中位数,可以中一分为二,将两边分别放在大小顶堆中;也可用TreeMap进行桶排,结合中位双指针来快速获取中位数。

一、数据流的中位数

image.png

二、双堆 & 桶

1、大小顶堆

// 数据流的中位数。
// 做不来,看来官方题解的提示部分,马上积极思考,双队列如何做。
// 做不来,先看提示 -> 独立思考 -> 学习别人沉淀的东西(尤其是优秀楼主,有独立思考,不看优秀者的总结,得不到升华;无独立思考,看了也是白看,get不到。)
public class MedianFinder {
    // 采用双优先队列,一个大顶堆(存左边)+一个小顶堆(存右边),保持2个队列size差1即可。
    PriorityQueue<Integer> bigTopQue = new PriorityQueue<>((p, q) -> q - p);
    PriorityQueue<Integer> smallTopQue = new PriorityQueue<>();

    public MedianFinder() {

    }

    public void addNum(int num) {
        if (bigTopQue.isEmpty()) bigTopQue.add(num);
        else if (smallTopQue.isEmpty()) smallTopQue.add(num);
        else {
            int[] nums = new int[3];
            // 将3数排序。
            sort(nums, bigTopQue.poll(), num, smallTopQue.poll());
            // 分发最大和最小数。
            bigTopQue.offer(nums[0]);
            smallTopQue.offer(nums[2]);
            // 根据情况分发中间的数。
            if (bigTopQue.size() <= smallTopQue.size()) bigTopQue.offer(nums[1]);
            else smallTopQue.offer(nums[1]);
        }
    }

    private void sort(int[] nums, int m, int num, int n) {
        nums[0] = m;
        nums[1] = num;
        nums[2] = n;

        Arrays.sort(nums);
    }

    public double findMedian() {
        int total = bigTopQue.size() + smallTopQue.size();
        int value = smallTopQue.peek();

        if ((total & 1) == 1) return value;

        return (bigTopQue.peek() + value) / 2d;
    }
}

2、TreeMap+双指针

// 类似于计数排序。非重复N个桶。
class MedianFinder2 {
    // 用TreeMap来排序,用双指针来记录左右中位数。
    TreeMap<Integer, Integer> tm = new TreeMap<>();
    // 第一位指向中位数,第二位指向重复中位数的第N个。需要根据N的大小来判定是否需要ceiling/flooring!
    int[] left = new int[2];
    int[] right = new int[2];

    public MedianFinder2() {

    }

    public void addNum(int num) {
        tm.put(num, tm.getOrDefault(num, 0) + 1);
        // 当容器为空时,将左右指针指向同一个元素,且是该元素的第一个值(重复情况)。
        int size = tm.size() - 1;
        if (size == 0) {
            left[0] = right[0] = num;
            left[1] = right[1] = 1;
        }
        // 当容器个数为偶数时,再根据num和left[0]/right[0]的大小来完成left/right数组的修改。
        else if ((size & 1) == 0) {
            // num在中间,则left需前进。
            if (left[0] < num && num < right[0]) increase(left);
                // num在left的左边,right-- && left应该和right指向同一个位置。
            else if (num <= left[0]) {
                decrease(right);
                // 防止left[0] == num,将num插在left[0]指向数的后面。
                System.arraycopy(right, 0, left, 0, 2);
            }
            // num 在right的右边,将left++即可。
            else increase(left);// 当right[0] == num时,插右边无影响,所以无需arraycopy
        }
        // 当容器个数为奇数时,两指针指向同一个元素,即right == left。
        else {
            // 当num < left[0]时,left--即可。
            if (num < left[0]) decrease(left);
                // 当num >= left[0]时,即大于right[0]时,直接right++即可。
            else increase(right);
        }
    }

    private void decrease(int[] iter) {
        if (--iter[1] == 0) {
            iter[0] = tm.floorKey(iter[0] - 1);
            iter[1] = tm.get(iter[0]);
        }
    }

    private void increase(int[] iter) {

        if (++iter[1] > tm.get(iter[0])) {
            iter[0] = tm.ceilingKey(iter[0] + 1);
            iter[1] = 1;
        }
    }


    public double findMedian() {
        return (left[0] + right[0]) / 2.0;
    }
}

总结

1)仔细分析问题,利用合适的数据结构来解决问题。除此之外,大小顶堆可快速获取最大/最小元素,对数级别。

2)TreeMap,以Key为比较对象,除了能O(1)的获取指定元素,还能配合ceiling/flooring来获取指定元素的左右元素。

3)做不来,看来官方题解的提示部分,马上积极思考,双队列如何做,养成积极思考的习惯的同时,也不忘学习他人,节约时间。

4)做不来,先看提示 -> 独立思考 -> 学习别人沉淀的东西(尤其是优秀楼主,有独立思考,不看优秀者的总结,得不到升华;无独立思考,看了也是白看,get不到。)

参考文献

[1] LeetCode 数据流的中位数