python最大堆heapq_python heapq 堆 求最大N和最小N个元素

这个模块(build-in)实现了一个堆的数据结构,完美的解决了Top-K问题,以后解决Top-K问题的时候,直接把这个模块拿来用就可以了

注意,默认的heap是一个小顶堆!

heapq模块提供了如下几个函数:

heapq.heappush(heap, item) 把item添加到heap中(heap是一个列表)

heapq.heappop(heap) 把堆顶元素弹出,返回的就是堆顶

heapq.heappushpop(heap, item) 先把item加入到堆中,然后再pop,比heappush()再heappop()要快得多

heapq.heapreplace(heap, item) 先pop,然后再把item加入到堆中,比heappop()再heappush()要快得多

heapq.heapify(x) 将列表x进行堆调整,默认的是小顶堆

heapq.merge(*iterables) 将多个列表合并,并进行堆调整,返回的是合并后的列表的迭代器

heapq.nlargest(n, iterable, key=None) 返回最大的n个元素(Top-K问题)

heapq.nsmallest(n, iterable, key=None) 返回最小的n个元素(Top-K问题)

求最大N个元素的值和下标

import numpy as np

import heapq

a = np.random.randint(0, 100, 10)

print(a)

print(heapq.nlargest(3, a))

print(heapq.nsmallest(3, a))

idx = heapq.nlargest(3, range(len(a)), a.take)

print('idx ', idx)

print(a[idx])

[57 82 47 31 87 66 40 64 95 35]

[95, 87, 82]

[31, 35, 40]

idx [8, 4, 1]

[95 87 82]

复杂结构时的使用

import numpy as np

import heapq

portfolio = [

{'name': 'IBM', 'shares': 100, 'price': 91.1},

{'name': 'AAPL', 'shares': 50, 'price': 543.22},

{'name': 'FB', 'shares': 200, 'price': 21.09},

{'name': 'HPQ', 'shares': 35, 'price': 31.75},

{'name': 'YHOO', 'shares': 45, 'price': 16.35},

{'name': 'ACME', 'shares': 75, 'price': 115.65}

]

cheap = heapq.nsmallest(3, portfolio, key=lambda s: s['price'])

print(cheap)

[{'name': 'YHOO', 'shares': 45, 'price': 16.35},

{'name': 'FB', 'shares': 200, 'price': 21.09},

{'name': 'HPQ', 'shares': 35, 'price': 31.75}]

从例子中可以看出,key匹配了portfolio中关键字为‘price’的一行。

到此为止,关于如何应用heapq来求Top N问题,相比通过上面的例子讲解,已经较为熟悉了。现在有几个需要注意的地方:

1)heapq.heapify(iterable):可以将一个列表转换成heapq

2)在Top N问题中,如果N=1,则直接用max(iterable)/min(iterable)即可。

3)如果N很大,接近集合元素,则为了提高效率,采用sort+切片的方式会更好,如:

求最大的N个元素:sorted(iterable, key=key, reverse=True)[:N]

求最小的N个元素:sorted(iterable, key=key)[:N]

你可能感兴趣的:(python最大堆heapq)