Python数据结构之堆并实现从海量数据中寻找最大的K个

1.堆其实是完全二叉树，有最大堆和最小堆

最大堆：对每个非叶子结点V，V的值都比它的两个孩子大

最大堆支持每次 pop 操作获取最大的元素，最小堆获取最小元素

常见问题：用堆来完成 topk 问题，从海量数字中寻找最大的 K 个

import heapq

class TopK:
	'''
	获取大量元素 topk 大个元素，固定内存
	思路：
	1.先放入元素前 K 个建立最小堆
	2.迭代剩余元素：
		如果当前元素小于堆顶元素，跳过该元素(肯定不是前 K 大)
		否则替换堆顶元素为当前元素，并重新调整堆
	'''
	
	def __init__(self, iterable, k):
		self.minheap = []         # 定义最小堆
		self.capacity = k         # 最小堆容量为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:
				# 返回并且pop堆顶最小值，推入新的 val 值并调整堆
				heapq.heapreplace(self.minheap, val)
		else:
			# 前面k个元素直接放入minheap
			heapq.heappush(self.minheap, val)
	
	def get_topk(self):
		'''获取topk的值'''
		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())

test()

输出结果：

    [990, 991, 993, 992, 997, 996, 994, 995, 998, 999]