求資料流中的中位數問題

Grey Zeng發表於2022-07-14

求資料流中的中位數問題

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

更多

演算法和資料結構筆記

相關文章