101.数据流的中位数

题目

中位数是有序整数列表中的中间值。如果列表的大小是偶数，则没有中间值，中位数是两个中间值的平均值。

例如 arr = [2,3,4] 的中位数是 3 。
例如 arr = [2,3] 的中位数是 (2 + 3) / 2 = 2.5 。

实现 MedianFinder 类:

MedianFinder() 初始化 MedianFinder 对象。
void addNum(int num) 将数据流中的整数 num 添加到数据结构中。
double findMedian() 返回到目前为止所有元素的中位数。与实际答案相差 10-5 以内的答案将被接受。

示例 1：

输入

["MedianFinder", "addNum", "addNum", "findMedian", "addNum", "findMedian"]         
[[], [1], [2], [], [3], []]

输出
[null, null, null, 1.5, null, 2.0]

解释:

MedianFinder medianFinder = new MedianFinder();
medianFinder.addNum(1);    // arr = [1]
medianFinder.addNum(2);    // arr = [1, 2]
medianFinder.findMedian(); // 返回 1.5 ((1 + 2) / 2)
medianFinder.addNum(3);    // arr[1, 2, 3]
medianFinder.findMedian(); // return 2.0

思路

我们要求的是：中间的数字。

例如: 1, 3, 5, 7, 9，要门其实不关心1和9，我们只关心中间，我们只要把数据分成两半，那就是两个堆。

想象你把所有已经添加的数字排好序，然后从中间切开:

左半部分: 较小的那一半(如果总数是奇数，就包含中位数)
右半部分: 较大的那一半

我们希望:

左半部分的最大值 <= 右半部分的最小值
左半部分的元素个数要么等于右半部分，要么比右半部分多一个。

那用什么数据结构？

左半部分: 用一个最大堆(Max Heap)，堆顶是左半部分最大的数。
右半部分: 用一个最小堆(Min Heap)，堆顶是右半部分最小的数。

在 Java/Kotlin 中，默认的 PriorityQueue 是最小堆，要实现最大堆可以用 Collections.reverseOrder() 或 compareBy { -it }。

这样中位数就可以直接从两个堆的堆顶得到。

要解决的问题

每次调用addNum(num)时，新数字来了，我们应该把它放到左边的大顶堆(left)，还是右边的小顶堆(right)。如果num <= left.peek()，说明这个数字比较小，应该放左边，否则放右边。
但是还有一个非常重要的问题：数据量必须平衡

我们不能让:

left：1
right：3 4 5 6 7 8 9

这样就失去意义了，我们要求两个堆的大小：最多相差1。

所以每次加入数字之后，我们都要平衡
规则可以简单记成:

如果 left.size > right.size + 1
    把 left 最大的拿出来
    放到 right

如果 right.size > left.size + 1
    把 right 最小的拿出来
    放到 left

这一步就是：两个堆互相搬元素。

获取中位数

如果总数为奇数，中位数 = 左堆堆顶或右堆对顶，看哪个数量多1
如果总数为偶数，中位数 = (左堆堆顶 + 右堆堆顶) / 2.0

import java.util.PriorityQueue

class MedianFinder {

    // 左边：大顶堆
    // 保存较小的一半数字
    private val left = PriorityQueue<Int>(
        compareByDescending { it }
    )

    // 右边：小顶堆
    // 保存较大的一半数字
    private val right = PriorityQueue<Int>()

    fun addNum(num: Int) {
        // 如果左边为空，或者 num 比左边最大值小
        // 放到左边
        if (left.isEmpty() || num <= left.peek()) {
            left.offer(num)
        } else {
            // 否则放右边
            right.offer(num)
        }

        // 保持两个堆大小最多相差 1
        if (left.size > right.size + 1) {
            right.offer(left.poll())
        }

        if (right.size > left.size + 1) {
            left.offer(right.poll())
        }
    }

    fun findMedian(): Double {
        // 左边多一个
        if (left.size > right.size) {
            return left.peek().toDouble()
        }

        // 右边多一个
        if (right.size > left.size) {
            return right.peek().toDouble()
        }

        // 两边数量相同
        return (left.peek().toDouble() + right.peek().toDouble()) / 2.0
    }
}

复杂度

时间复杂度: O(log n)（堆的插入和删除）
空间复杂度: O(1)

下一题

邮箱：charon.chui@gmail.com
Good Luck!