求数据流中的中位数问题 - Grey Zeng

2 years ago

source link: https://www.cnblogs.com/greyzeng/p/16479520.html
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.

求数据流中的中位数问题

作者：Grey

原文地址: 求数据流中的中位数问题

题目链接#

LeetCode 295. Find Median from Data Stream

主要思路#

要得到数据流中的中位数，在偶数的情况下，要得到上下中位数求平均，在奇数的状态下，要得到中间位置的数，这里最关键的问题就是：如何快速找到中间位置（而且是动态的）。

我们可以准备两个堆，一个大根堆，一个小根堆，两个堆分别维持在一半左右的元素，且让这两个堆的堆顶始终保持中间位置的元素。这样就可以快速得到中位数需要的值了。具体操作如下：

第一个元素永远是先进大根堆。

接下来的元素，按如下规则进入：

如果接下来的元素比大根堆堆顶元素小，进大根堆，否则进入小根堆。

如果两个堆的大小差值已经达到2了，说明元素要向着一侧堆倾斜，这个时候，为了维持两个堆的平衡（即：始终可以拿到中位数需要的信息），从数量多的堆拿出一个放到数量少的堆，这样就让两个堆始终保持差值小于等于1。

此时，如果要得到中位数，通过如下规则就可以得到：

如果大根堆和小根堆数量一样，说明原始数据流是偶数个，那么直接拿出大根堆和小根堆的堆顶元素求和再除以2就是了。

如果大根堆和小根堆数量不一样（在这一步，只能是差1），那么就取多的那个堆的堆顶即为中位数。

完整代码如下

import java.util.Comparator;
import java.util.PriorityQueue;

 class MedianFinder {

        private final PriorityQueue<Integer> minHeap;
        private final PriorityQueue<Integer> maxHeap;

        public MedianFinder() {
            minHeap = new PriorityQueue<>(Comparator.comparingInt((Integer o) -> o));
            maxHeap = new PriorityQueue<>((o1, o2) -> o2 - o1);
        }

        public void addNum(int num) {
            if (maxHeap.isEmpty()) {
                maxHeap.add(num);
            } else {
                if (maxHeap.peek() >= num) {
                    maxHeap.add(num);
                } else {
                    minHeap.add(num);
                }
            }
            int maxHeapSize = maxHeap.size();
            int minHeapSize = minHeap.size();
            if (maxHeapSize - minHeapSize == 2) {
                minHeap.add(maxHeap.poll());
            } else if (minHeapSize - maxHeapSize == 2) {
                maxHeap.add(minHeap.poll());
            }
        }

        public double findMedian() {
            if (maxHeap.size() == minHeap.size()) {
                return (maxHeap.peek() + minHeap.peek()) / 2d;
            }
            if (maxHeap.size() > minHeap.size()) {
                return maxHeap.peek();
            }
            return minHeap.peek();
        }
}