Python中的堆实现：heapq 模块——利用堆结构实现快速访问数据流中的中位数

堆结构

堆结构是一种优先队列，可以以任意顺序添加对象，并随时查找或删除最小（大）的元素，或者查找和删除前 K 个最小（大）元素。相比于列表方法min() / max()，这样做的效率要高得多。

堆结构是一种特殊的完全二叉树（除了叶子节点层外，其余层节点数均达到最大值，而叶子节点层所有节点都集中在左侧）。根节点的值不大于（小于）其子节点的值，并且子节点也服从这种特性。根节点值不大于子节点的堆称为小根堆，根节点的值不小于子节点的堆称为大根堆。如图左为小根堆，图右为大根堆。

Python中 heapq 模块

Python 中给出了小根堆的辅助实现库函数 heapq 模块（其中q表示队列，方便记忆），大根堆可以通过将数值取反来实现。heapq 模块只是维护堆结构的操作函数，因此要用列表 list 来表示堆结构本身。

 from heapq import *          # 使用前首先导入heapq模块
 
 heap = [3, 5, 1, 7, 8, 9, 0, 2, 4, 6]
 heapq.heapify(heap)          # 让列表 heap具有堆属性
 print(heap)                  # [0, 2, 1, 4, 6, 9, 3, 7, 5, 8] heap[0]值为堆最小值
 
 heappush(heap, -1)           # 向堆中添加元素，也可以用于创建堆
 print(heap)                  # [-1, 0, 1, 4, 2, 9, 3, 7, 5, 8, 6] heap[0]依旧为最小值

 min_value = heappop(heap)    # 弹出堆中最小值
 print(heap, min_value)       # [0, 2, 1, 4, 6, 9, 3, 7, 5, 8] -1

 min_value = heapreplace(heap, 4)       # 弹出堆中最小值，并将值4加入堆中并维持堆结构。
 print(heap, min_value)                 # [1, 2, 3, 4, 6, 9, 4, 7, 5, 8] 0
 
 nList = nlargest(3, heap)              # 返回 heap中前3大的元素
 print(nList)                           # [9, 8, 7]
 nList = nsmallest(3, heap)             # 返回 heap中最小的3个元素
 print(nList)                           # [1, 2, 3]
 
 small_num = heappushpop(heap, -2)      # 将数值-2压入堆 heap中，并返回弹出堆的最小值
 print(heap, small_num)                 # [1, 2, 3, 4, 6, 9, 4, 7, 5, 8] -2

总结：

函数	描述
heapify(heap)	以线性时间将一个列表 heap 转化为堆
heappush(heap, num)	将元素 num 压入堆 heap 中
heappop(heap)	从对 heap 中弹出最小元素，并返回
heap[0]	查看堆 heap 中最小值，不弹出
heapreplace(heap, num)	弹出并返回最小值后，将 num 压入堆中并维持堆结构，比先调用 heappop(), 再调用heappush() 效率高
nlargest(n, heap)	返回堆 heap 中 n 个最大元素
nsmallest(n, heap)	返回堆 heap 中 n 个最小元素
heappushpop(heap, num)	将值 num 插入到堆 heap 中后，再返回并弹出最小值。和先调用 heappush(), 再调用 heappop() 相比效率更高

为 nlargest() 和 nsmallest() 函数指定 key 值进行比较。两个函数共有三个参数，原型为：nlargest(n , iterbale, key=None) 和 nsmallest(n , iterbale, key=None)。key 值表示用列表元素的某个属性和函数作为关键字进行判断。

L = [('a', 1), ('b', 2), ('c', 3), ('d', 0)]
nList = nlargest(2, L, key = lambda x:x[1])  
# key=lambda x:x[1] 表示以前面对象中第二维数据进行排列，更一般的形式为：key=lambda 变量:变量[维数]
print(nList)   # [('c', 3), ('b', 2)]

利用堆结构快速访问数据流中中位数

对应LeetCode295题：https://leetcode-cn.com/problems/find-median-from-data-stream/
题目描述：如何得到一个数据流中的中位数？如果从数据流中读出奇数个数值，那么中位数就是所有数值排序之后位于中间的数值。如果从数据流中读出偶数个数值，那么中位数就是所有数值排序之后中间两个数的平均值。
代码实现：

import heapq

class MedianFinder:

    def __init__(self):
        self.smallheap = []       # 初始化一个小根堆
        self.bigheap = []         # 初始化一个大根堆，通过对元素取反实现

    def addNum(self, num: int) -> None:   # 添加元素操作，使大根堆中元素始终小于小根堆，两个堆中元素数目差1，或相等
        if len(self.smallheap) == len(self.bigheap):     # 当前数据中含偶数个数据
            # 将新加入元素取反，加入大根堆中并返回大根堆中最小值并取反，也就是实际元素最大值。然后将这个值压入小根堆中。
            heapq.heappush(self.smallheap, -heapq.heappushpop(self.bigheap, -num))
        else:    # 当前数据含奇数个数据
            # 将新数据压入小根堆中，并返回最小值后取反压入最大堆
            heapq.heappush(self.bigheap, -heapq.heappushpop(self.smallheap, num))

    def findMedian(self) -> float:  # 返回数据的中位数
        if len(self.smallheap) == len(self.bigheap):  # 当含有偶数个数据时，中位数为小根堆最小值和大根堆最大值的平均值
            return (self.smallheap[0] - self.bigheap[0])/2
        else:     # 当含有奇数个数据时，中位数为小跟堆最小值。
            return self.smallheap[0]

原创不易，如果感觉有所帮助，请点赞支持！感谢！