python利用优先队列取最大K个元素

利用priorityqueue取最大的K个元素

在进行稍微大一点的数据排序时, 有时并不需要全部数据, 只需要前面的最大k个数, 如果用sorted([…])[::-1][:k]实现效率会低很多, 为了优化执行速度, 可以采用优先队列priorityqueue实现求最大的K个元素. 同时空间也可以得到明显优化

import queue
import random

K = 12
q = queue.PriorityQueue()
tlist = []

for ith in range(100):
    q.put((ith, ith + random.randrange(100)))
    if q.qsize() > K: q.get() # pop least item
        
while not q.empty():
    tnum = q.get()
    tlist.append(tnum)

tlist.sort(reverse=True)
print(tlist)

效果

[(99, 187), (98, 152), (97, 184), (96, 136), 
(95, 105), (94, 119), (93, 109), (92, 169), 
(91, 149), (90, 143), (89, 91), (88, 92)]

priorityqueue是在queue包里, 而queue里priorityqueue的实现是调用的headq里小顶堆的实现, 所以没有提供key参数改变优先队列的排序比较规则, 就只有从小到大升序排列, 所以为了取前K大元素, 需要在队列规模大于K时弹出队首元素, 维持最大的K个在队列里, 最后再用一个list承接队列里的元素, 再降序排序一次就得到目标结果.