算尽天下|算法|题解

[算尽天下系列第3期]LeetCode·295·数据流的中位数

凝神长老 · 2月15日 · 2020年 · · · · 326次已读

题目描述

中位数是有序列表中间的数。如果列表长度是偶数,中位数则是中间两个数的平均值。

例如,

[2,3,4] 的中位数是 $3$

[2,3] 的中位数是 $(2 + 3) / 2 = 2.5$

设计一个支持以下两种操作的数据结构:

void addNum(int num) – 从数据流中添加一个整数到数据结构中。
double findMedian() – 返回目前所有元素的中位数。

来源:力扣(LeetCode
链接:https://leetcode-cn.com/problems/find-median-from-data-stream\

样例

样例

addNum(1)
addNum(2)
findMedian() -> 1.5
addNum(3) 
findMedian() -> 2

算法与数据结构

题解

一种显然的做法是,将数字存储 ArrayList 中,每次需要输出中位数时,对其进行排序,并输出中间值。

对于一组静态数据,中位数是固定的,我们可以先排序,第 $\frac{n}{2}$ 个数据就是中位数。不论询问多少次中位数,直接返回这个固定的值就好了。所以,尽管排序的代价比较大,但是边际成本会很小。

但是,如果我们面对的是动态数据集合,中位数在不停地变动,如果再用先排序的方法,每次询问中位数的时候,都要先进行排序,那效率就不高了。

一种显然的做法是使用优先级队列,对于一个优先级队列,每次遇到 addNum() 就往里面加入新的数据,优先级队列会自动帮我们从小到大排序,而每次遇到 findMedian() 的时候,就输出 pq.size() / 2 位置上的值(如果 size 为奇数)或者 pq.size() / 2pq.size() / 2 + 1 这两个位置的平均值(如果 size 为偶数)。

然而,对于一般的优先级队列,只能访问队首的元素,并不提供随机访问的接口,这时候就只能手动维护两个堆了。

我们需要维护两个堆,一个大顶堆,一个小顶堆。大顶堆中存储前半部分数据,小顶堆中存储后半部分数据,且小顶堆中的数据都大于大顶堆中的数据。

也就是说,如果有 $n$ 个数据,$n$ 是偶数,我们从小到大排序,那前 $\frac{n}{2}$ 个数据存储在大顶堆中,后 $\frac{n}{2}$ 个数据存储在小顶堆中。如果 $n$ 是奇数,情况是类似的,大顶堆就存储 $\frac{n}{2} + 1$ 个数据,小顶堆中就存储 $\frac{n}{2}$ 个数据。

在寻找中位数的时候,如果 $n$ 是奇数,那么中位数就是大顶堆的堆顶元素,如果 $n$ 是偶数,那么中位数就是大顶堆和小顶堆各自堆顶元素的平均值。这个结论是显然成立的。

于是,findMedian() 就很容易得到了。

if (max.size() == min.size()) {
    // n 是偶数
    return 0.5 * (max.peek() + min.peek());
} else {
    // n 是奇数
    return max.peek();
}

这样,我们所有的工作就集中在了堆的调整上,如何在插入新数据的时候,依然保证上面的约定条件成立呢?

如果新加入的数据小于等于大顶堆的堆顶元素,我们就将这个新数据插入到大顶堆;否则,我们就将这个新数据插入到小顶堆。

if (num < max.peek()) {
    max.add(num);
} else {
    min.add(num);
}

这个时候就有可能出现两个堆中的数据个数不符合前面约定的情况:如果 $n$ 是偶数,两个堆中的数据个数都是 $\frac{n}{2}$;如果 $n$ 是奇数,大顶堆有 $\frac{n}{2} + 1$ 个数据,小顶堆有 $\frac{n}{2}$ 个数据。
这个时候,我们可以从一个堆中不停地将堆顶元素移动到另一个堆,通过这样的调整,来让两个堆中的数据满足上面的约定。

我们可以写一个循环,当大顶堆的元素个数小于小顶堆的元素个数时,将小顶堆的堆顶元素出堆,并加入大顶堆;反之,则将大顶堆的堆顶元素出堆,加入小顶堆。这个循环的终止条件是,两个堆的元素个数相等($n$ 为偶数),或者大顶堆比小顶堆多 $1$ 个元素($n$ 为奇数)。

while (true) {
    int minSize = min.size();
    int maxSize = max.size();
    if (minSize == maxSize || minSize + 1 == maxSize) {
        break;
    } else {
        if (minSize < maxSize) {
            min.add(max.poll());
        } else {
            max.add(min.poll());
        }
    }
}

最后检查一下 Corner Case。

对于本题来说,有可能出现错误的地方,就是当堆为空的时候,访问 peek() 的结果不能交给 int

findMedian() 中,我们可以很自信地说出,这样的情况不可能存在。

因为既然是找中位数,则至少有 $1$ 个元素,而大顶堆的元素个数一定大于或者等于小顶堆的元素个数,因此,大顶堆元素个数不可能为 $0$,大顶堆不可能为空。

既然大顶堆不可能为空,那么,若 $n$ 为奇数,返回大顶堆的堆顶元素,就不可能遇到 null,而若 $n$ 为偶数,说明大顶堆和小顶堆的元素个数相等,那么小顶堆自然也不可能为空,因此 peek() 也是安全的,可以正常计算两个数的均值。

public double findMedian() {
    assert !max.isEmpty();
    if (max.size() == min.size()) {
        return 0.5 * (max.peek() + min.peek());
    } else {
        return max.peek();
    }
}

另外还有一个地方用到了 peek(),那就是在插入的地方。

上面说到,插入新的数据的时候,永远是与大顶堆的堆顶元素比较,而大顶堆不可能为空,所以也是安全的?

等等,那么,第一个元素怎么办?

简单做一下特判就可以了。

if (max.isEmpty() || num < max.peek()) {}

完整代码

class MedianFinder {

    private static PriorityQueue<Integer> min;
    private static PriorityQueue<Integer> max;

    public MedianFinder() {
        min = new PriorityQueue<>();
        max = new PriorityQueue<>(new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                return o2-o1;
            }
        });
    }

    public void addNum(int num) {
        if (max.isEmpty() || num < max.peek()) {
            max.add(num);
        } else {
            min.add(num);
        }
        while (true) {
            int minSize = min.size();
            int maxSize = max.size();
            if (minSize == maxSize || minSize + 1 == maxSize) {
                break;
            } else {
                if (minSize < maxSize) {
                    min.add(max.poll());
                } else {
                    max.add(min.poll());
                }
            }
        }
    }

    public double findMedian() {
        assert !max.isEmpty();
        if (max.size() == min.size()) {
            return 0.5 * (max.peek() + min.peek());
        } else {
            return max.peek();
        }
    }
}

备注

时间复杂度: $O(5 \cdot \log n) + O(1) \approx O(\log n)$。

最坏情况下,从顶部有三个堆插入和两个堆删除。每一个都需要花费 $O(\log n)$ 时间。

找到平均值需要持续的 $O(1)$ 时间,因为可以直接访问堆的顶部。

空间复杂度:$O(n)$ 用于在容器中保存输入的线性空间。

另外,在上文中提到了

if (pq.size() % 2 == 0) {
    return 0.5 * (pq.get(pq.size()) + pq.get(pq.size() + 1));
} else {
    return pq.get(pq.size() / 2);
}

自平衡二进制搜索树(如 AVL 树)具有一些非常有趣的特性。它们将树的高度保持在对数范围内。因此,插入新元素具有相当好的时间性能。中值总是缠绕在树根或它的一个子树上。大多数语言实现模拟这种行为的是 multiset 类。唯一的问题是跟踪中值元素。

我们保持两个指针:一个用于中位数较低的元素,另一个用于中位数较高的元素。当元素总数为奇数时,两个指针都指向同一个中值元素(因为在本例中只有一个中值)。当元素数为偶数时,指针指向两个连续的元素,其平均值是输入的代表中位数。

0 0 投票
Article Rating
订阅评论动态
提醒
guest
0 评论
行内反馈
查看所有评论