Python heapq实现大顶堆_python heapq大顶堆

Python的内置标准库heapq实现的是小顶堆，要想利用heapq实现大顶堆有两种思路：

• 将数据取反进行push（此方法只针对数值型数据）
• 重载__lt__函数（针对自定义数据类型使用）

将数据取反进行push

该方法很简单，如其名，将数据取反后push到堆中，pop的时候再取反即可。示例代码如下：

# heapq实现大顶堆
import numpy as np
from heapq import heappop as pop
from heapq import heappush as push

if __name__ == '__main__':
    N = 30
    a = np.random.randint(1,500,N)

    heap = []

    for i in a:
        push(heap, -i)

    ans = [-pop(heap) for _ in range(N)]
    
    print(ans)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

运行结果如下

[483, 395, 375, 368, 359, 355, 345, 334, 330, 327, 325, 321, 308, 291, 273, 249, 244, 229, 201, 192, 173, 152, 138, 128, 54, 45, 26, 17, 10, 8]

重载__lt__函数

针对我们自定义的数据结构，可以采用重载__lt__函数的方式实现大顶堆。示例代码如下：

# 自定义数据结构heapq实现大顶堆
import numpy as np
from heapq import heappop as pop
from heapq import heappush as push

class  my_data:
    def __init__(self, a, b) -> None:
        self.a = chr(a) # 为演示非数值型数据，这里定义字符类型数据
        self.b = b
    def __lt__(self, data):
        return self.a > data.a

    
if __name__ == '__main__':

    # print(dir(my_data))
    N = 50
    test_data = []

    for i in range(N):
        test_data.append(my_data(np.random.randint(65, 90), np.random.randint(500)))

    heap = []
    
    for i in test_data:
        push(heap, i)

    # 展示默认排序结果
    ans = [(pop(heap)) for _ in range(N)]
    ans = [(i.a, i.b) for i in ans]
    print(ans)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

运行结果如下

[(‘X’, 225), (‘X’, 475), (‘W’, 211), (‘V’, 41), (‘V’, 269), (‘V’, 475), (‘U’, 429), (‘U’, 361), (‘U’, 370), (‘T’, 58), (‘T’, 167), (‘T’, 236), (‘S’, 328), (‘S’, 350), (‘S’, 473), (‘S’, 102), (‘R’, 489), (‘R’, 214), (‘Q’, 165), (‘Q’, 270), (‘P’, 166), (‘P’, 54), (‘P’, 255), (‘O’, 38), (‘M’, 64), (‘M’, 70), (‘L’, 367), (‘L’, 311), (‘L’, 334), (‘L’, 489), (‘J’, 75), (‘I’, 103), (‘H’, 499), (‘G’, 301), (‘G’, 294), (‘G’, 283), (‘G’, 116), (‘F’, 96), (‘F’, 186), (‘F’, 205), (‘E’, 96), (‘E’, 364), (‘D’, 407), (‘C’, 494), (‘C’, 335), (‘B’, 43), (‘B’, 193), (‘B’, 271), (‘B’, 414), (‘A’, 76)]

这样就实现了自定义数据的大顶堆。