大根堆与小根堆(java,c++,python)

大根堆与小根堆(java,c++,python)

java语言

Java 使用 PriorityQueue<>((x, y) -> (y - x)) 可方便实现大顶堆。

    PriorityQueue<Integer> l = new PriorityQueue<>((a,b)->b-a);
    PriorityQueue<Integer> r = new PriorityQueue<>((a,b)->a-b);

l 为大根堆，r 为小根堆

c++语言

    // 大顶堆, 存储较小一半的值，降序队列，堆顶为最大值。
    priority_queue<int, vector<int>, less<int>> l;

    // 小顶堆, 存储较大一半的值，升序队列，堆顶为最小值。
    priority_queue<int, vector<int>, greater<int>> r;

l 为大根堆，r 为小根堆

python语言

Python 中 heapq 模块是小顶堆。实现 大顶堆 方法： 小顶堆的插入和弹出操作均将元素 取反 即可。

from heapq import *

class MedianFinder:
    def __init__(self):
        self.A = [] # 小顶堆，保存较大的一半
        self.B = [] # 大顶堆，保存较小的一半

    def addNum(self, num: int) -> None:
        if len(self.A) != len(self.B):
            heappush(self.A, num)
            heappush(self.B, -heappop(self.A))
        else:
            heappush(self.B, -num)
            heappush(self.A, -heappop(self.B))

    def findMedian(self) -> float:
        return self.A[0] if len(self.A) != len(self.B) else (self.A[0] - self.B[0]) / 2.0

参考

面试题41. 数据流中的中位数（优先队列 / 堆，清晰图解）
优先队列的用法（priority_queue）