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

前言

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

一、数据流的中位数

二、双堆 & 桶

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;
    }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

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
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62

总结

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

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

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

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

参考文献

[1] LeetCode 数据流的中位数