Python heapq 堆数据结构

问题 : 我们想在某个集合中找出最大或最小的N个元素

heapq

说明 : heapq 模块中有两个函数

• heapq.nlargest(n, iterable[, key]) : 从迭代器对象iterable中返回前n个最大的元素列表,其中关键字参数key用于匹配是字典对象的iterable，用于更复杂的数据结构中。

• heapq.nsmallest(n, iterable[, key]) : 从迭代器对象iterable中返回前n个最小的元素列表，其中关键字参数key用于匹配是字典对象的iterable，用于更复杂的数据结构中。

import heapq
nums = [1,8,2,23,7,-4,18,23,42,37,2]
print(heapq.nlargest(3, nums))  #  [42, 37, 23]
print(heapq.nsmallest(3, nums))  #  [-4, 1, 2]

在两个函数中都可以接受参数key,从而允许它们在工作在更加复杂的数据结构之上

# 根据字典的某个键值对进行排列
portfolio = [
    {"name": "IBM", "shares": 100, "price": 91.1},
    {"name": "AAPL", "shares": 50, "price": 533.22},
    {"name": "FB", "shares": 200, "price": 21.09},
    {"name": "HPQ", "shares": 35, "price": 31.75},
    {"name": "YHOO", "shares": 45, "price": 16.35},
    {"name": "IBM", "shares": 75, "price": 115.65},
]

cheap = heapq.nsmallest(3,portfolio,key=lambda s :s["price"])
expensive = heapq.nlargest(3,portfolio,key=lambda s :s["price"])
print(cheap)
print(expensive)

问题 : 当从一个很大的集合中去找最大或者最小的 N 个元素, 且同集合中元素的总数目相比, N 很小,那么如何才能更高效率的找到 N 个元素的那?

首先我们需要将函数在底层将数据转化为列表,然后通过 heapq.heapify() 元素会以堆得顺序排列.当元素以堆排序的时候此时会有一个特性就是 heap[0] 总是最小的那个元素. 此外,接下来的元素可依次通过 heapq.heappop() 方法轻松找到. 该方法会将第一个元素(最小的弹出),这个依次取出第二个,如果我们想找到第3小的元素.

nums = [1, 8, 2, 23, 7, -4, 18, 23, 42, 37, 2]
heap = list(nums)    # 在底层中我们首先要将数据转化为列表,而且这一步是必不可少的
heapq.heapify(heap)  # 函数没有返回值的
print(heap)          # [-4, 2, 1, 23, 7, 2, 18, 23, 42, 37, 8]
# 获取第3个最小的元素
print(heapq.heappop(heap))  # -4 
print(heapq.heappop(heap))  # 1
print(heapq.heappop(heap))  # 2

总结 : 从一个结合中如何寻找最大最小的N个元素的方法选择

简单地想找到最小或者最大的元素(N=1),那么使用 min() 和 max() 会更加快.

当N的大小和集合的本身的大小差不多大,通常最快的方法是先对集合进行排序,然后做切片操作sorted(items)[:N]/sorted(items)[-N:].

nlargest()/nsmallest() 的实际实现方式会根据使用它们的方式而有所不同,可能会出现相应的一些优化措施(比如,当N的大小同输入大小很接近的时候,就会采用排序的算法)