使用heapq构建最小堆和最大堆

python默认的heapq模块只支持最小堆的构建。昨天在做一个关于数据流中使用O(1)时间复杂度求中位数的算法题的时候,用到了最大堆和最小堆。借助heapq模块实现的最小堆如下

class SmallHeap:
    def init(self):
        self.arr = list()
    def heap_insert(self, val):
        heapq.heappush(self.arr, val)
    def heapify(self):
        heapq.heapify(self.arr)
    def heap_pop(self):
        return heapq.heappop(self.arr)
    def get_top(self):
        if not self.arr:
            return
        return self.arr[0]

heapq中的poppush等操作会自动调动heapify方法调整堆。除了上述方法,还有一些常用方法 - heappushpop: 相当于先执行了heappush(heap,item),然后执行了heappop(heap) - heapreplace: 相当于先进行heappop(heap),再进行heappush(heap,item)操作 - nlargest/nsmallest: 求堆中最大/最小的k个元素 - merge: 输入是两个有序的数组,输出是数组合并后的有序的generator。注意,如果输入无序,那么输出也不会保证有序

上面是小根堆的相关操作。python的heapq不支持大根堆,在stackoverflow上看到了一个巧妙的实现:我们还是用小根堆来进行逻辑操作,在做push的时候,我们把最大数的相反数存进去,那么它的相反数就是最小数,仍然是堆顶元素,在访问堆顶的时候,再对它取反,就获取到了最大数。思路很是巧妙。下面是实现代码

class BigHeap:
    def init(self):
        self.arr = list()
    def heap_insert(self, val):
        heapq.heappush(self.arr, -val)
    def heapify(self):
        heapq.heapify(self.arr)
    def heap_pop(self):
        return -heapq.heappop(self.arr)
    def get_top(self):
        if not self.arr:
            return
        return -self.arr[0]

heapq针对的是基本类型的堆操作,如果是复杂对象的比较的话,还是需要我们自己实现堆的insertheapify等操作。