求数据流中的中位数问题

求数据流中的中位数问题

作者：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();
        }
}

折叠 

更多#

算法和数据结构笔记