TOPK 問題
描述
如從海量數(shù)字中尋找最大的 k 個愁溜,這類問題我們稱為 TOPK 問題疾嗅,通常使用堆來解決:
- 求前 k 大,用最小堆
- 求前 k 小冕象,用最大堆
例子
現(xiàn)有列表 [1, 2, 0, 3, 5]
, 求前 2 個大的元素代承。
如傳入列表和 k = 2,輸出 [3, 5]
渐扮。
思路
先放入元素前 k 個建立一個最小堆
-
迭代剩余元素:
如果當(dāng)前元素小于堆頂元素论悴,跳過該元素(肯定不是前 k 大)
否則替換堆頂元素為當(dāng)前元素,并重新調(diào)整堆
最后獲取 最小堆 中的值墓律,即為 topk
代碼如下
import heapq
class Topk:
"""獲取大量元素 topk 大個元素膀估,固定內(nèi)存
思路:
1. 先放入元素前 k 個建立一個最小堆
2. 迭代剩余元素:
如果當(dāng)前元素小于堆頂元素,跳過該元素(肯定不是前 k 大)
否則替換堆頂元素為當(dāng)前元素只锻,并重新調(diào)整堆
"""
def __init__(self, iterable, k):
self.minheap = []
self.capacity = k
self.iterable = iterable
def push(self, val):
if len(self.minheap) >= self.capacity:
min_val = self.minheap[0]
if val < min_val: # 當(dāng)然你可以直接 if val > min_val 操作玖像,這里我只是顯示指出跳過這個元素
pass
else:
heapq.heapreplace(self.minheap, val) # 返回并且 pop 堆頂最小值,推出新的 val 值并調(diào)整堆
else:
heapq.heappush(self.minheap, val) # 前面 k 個元素直接放入 minheap
def get_topk(self):
for val in self.iterable:
self.push(val)
return self.minheap
def test():
import random
i = list(range(1000)) # 這里可以是一個可迭代元素齐饮,節(jié)省內(nèi)存
random.shuffle(i)
_ = Topk(i, 10)
print(_.get_topk()) # [990, 992, 991, 993, 996, 997, 998, 994, 995, 999]
test()