数据流的中位数 -- 大小根堆组合使用

295. 数据流的中位数

Java 实现

import java.util.PriorityQueue;

class MedianFinder {
    private PriorityQueue<Integer> large;
    private PriorityQueue<Integer> small;

    public MedianFinder() {
        // 小顶堆
        large = new PriorityQueue<>();
        // 大顶堆
        small = new PriorityQueue<>((a, b) -> {
            return b - a;
        });
    }

    public double findMedian() {
        // 如果元素不一样多，多的那个堆的堆顶元素就是中位数
        if (large.size() < small.size()) {
            return small.peek();
        } else if (large.size() > small.size()) {
            return large.peek();
        }
        // 如果元素一样多，两个堆堆顶元素的平均数是中位数
        return (large.peek() + small.peek()) / 2.0;
    }

    public void addNum(int num) {
        if (small.size() >= large.size()) {
            small.offer(num);
            large.offer(small.poll());
        } else {
            large.offer(num);
            small.offer(large.poll());
        }
    }

    public static void main(String[] args) {
        MedianFinder medianFinder = new MedianFinder();
        medianFinder.addNum(1);    // arr = [1]
        medianFinder.addNum(2);    // arr = [1, 2]
        System.out.println(medianFinder.findMedian()); // 返回 1.5 ((1 + 2) / 2)
        medianFinder.addNum(3);    // arr[1, 2, 3]
        System.out.println(medianFinder.findMedian()); // return 2.0
    }
}

Python 实现

import heapq

class MedianFinder:
    def __init__(self):
        # 小顶堆
        self.large = []
        # 大顶堆【heapq默认小根堆，所有插入元素时，用负数】
        self.small = []
    
    def findMedian(self) -> float:
        # 如果元素不一样多，多的那个堆的堆顶元素就是中位数
        if len(self.large) < len(self.small):
            return self.small[0]
        elif len(self.large) > len(self.small):
            return self.large[0]
        # 如果元素一样多，两个堆堆顶元素的平均数是中位数
        return (self.large[0] - self.small[0]) / 2
    
    def addNum(self, num: int) -> None:
        if len(self.small) >= len(self.large):
            heapq.heappush(self.small, -num)
            heapq.heappush(self.large, -heapq.heappop(self.small))
        else:
            heapq.heappush(self.large, num)
            heapq.heappush(self.small, -heapq.heappop(self.large))

if __name__ == '__main__':
	medianFinder = MedianFinder()
	medianFinder.addNum(1)    # arr = [1]
	medianFinder.addNum(2)    # arr = [1, 2]
	print(medianFinder.findMedian()) # 返回 1.5 ((1 + 2) / 2)
	medianFinder.addNum(3)    # arr[1, 2, 3]
	print(medianFinder.findMedian()) # return 2.0