python heapq堆排序高阶用法

python heapq相关操作

heapq 的一些常用方法：

1. heapify(list)将列表转换为小根堆的数据结构。

2. heappush(heap, x)，将 x 加入堆中。

3. heappop(heap)，从堆中弹出最小的元素。

4. heapreplace(heap, x)，弹出最小的元素，并将 x 加入堆中。

5. heappushpop(heap, x)，先把 x 加入堆中，再弹出最小的元素。

6. heapq.nlargest(n, heap)，返回 heap 中前 n 个最大的元素。

7. heapq.nsmallest(n, heap)，返回 heap 中前 n 个最小的元素。

如何自定义堆排序?

重写__lt__方法

import heapq

a = range(10)

class Node:
    def __init__(self, v):
        self.v = v

    def __lt__(self, other):
        # 自定义lt便可实现大顶堆
        return self.v > other.v

h = []
for i in a:
    heapq.heappush(h, Node(i))
ans = heapq.nsmallest(len(h), h)
for i in ans:
    print(i.v)

a = [(1, 'asd', 2), (2, 'asd', 2), (1, 'asd', 3), (5, 'asd', 3), (1, 'fd', 3)]
# 对于a想要按照第一个数从大到小，第二个字符串从大到小，第三个数字从小到大排序
class Node:
    def __init__(self, v):
        self.a, self.b, self.c = v[0], v[1], v[2]

    def __lt__(self, other):
        if self.a != other.a:
            return self.a > other.a
        if self.b != other.b:
            return self.b > other.b
        return self.c < other.c
h = []
for i in a:
    heapq.heappush(h,Node(i))
ans = heapq.nsmallest(len(h),h)
for i in ans:
    print(i.a,i.b,i.c)
# 5 asd 3
# 2 asd 2
# 1 fd 3
# 1 asd 2
# 1 asd 3

利用itemgetter

from operator import itemgetter
import heapq

list1 = [1, 6, 4, 3, 9, 5]
list2 = ['12', 'a6', '4', 'c34', 'b9', '5']
list3 = [
    {'name': 'jim', 'age': 23, 'price': 500},
    {'name': 'mase', 'age': 23, 'price': 600},
    {'name': 'tom', 'age': 25, 'price': 2000},
    {'name': 'alice', 'age': 22, 'price': 300},
    {'name': 'rose', 'age': 21, 'price': 2400},
]

print(heapq.nlargest(len(list1), list1))
print(heapq.nlargest(len(list2), list2))
print(heapq.nlargest(len(list3), list3, key=itemgetter('age', 'price')))

print(heapq.nsmallest(len(list1), list1))
print(heapq.nsmallest(len(list2), list2))
print(heapq.nsmallest(len(list3), list3, key=itemgetter('age', 'price')))