TOPK 问题
描述
如从海量数字中寻找最大的 k 个，这类问题我们称为 TOPK 问题，通常使用堆来解决：
- 求前 k 大，用最小堆
- 求前 k 小，用最大堆
例子
现有列表 [1, 2, 0, 3, 5], 求前 2 个大的元素。
如传入列表和 k = 2，输出 [3, 5]。
思路
先放入元素前 k 个建立一个最小堆
-
迭代剩余元素：
如果当前元素小于堆顶元素，跳过该元素（肯定不是前 k 大）
否则替换堆顶元素为当前元素，并重新调整堆
最后获取 最小堆 中的值，即为 topk
image.png
代码如下
import heapq
class Topk:
"""获取大量元素 topk 大个元素，固定内存
思路：
1. 先放入元素前 k 个建立一个最小堆
2. 迭代剩余元素：
如果当前元素小于堆顶元素，跳过该元素（肯定不是前 k 大）
否则替换堆顶元素为当前元素，并重新调整堆
"""
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: # 当然你可以直接 if val > min_val 操作，这里我只是显示指出跳过这个元素
pass
else:
heapq.heapreplace(self.minheap, val) # 返回并且 pop 堆顶最小值，推出新的 val 值并调整堆
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)) # 这里可以是一个可迭代元素，节省内存
random.shuffle(i)
_ = Topk(i, 10)
print(_.get_topk()) # [990, 992, 991, 993, 996, 997, 998, 994, 995, 999]
test()
