Python 用最小堆獲取大量元素 topk 大個元素

import heapq class TopK: """ 獲取大量元素 topk 大個元素，固定內存 思路： 1. 先讓入元素前 k 個建立一個最小堆 2. 迭代剩餘元素： 如果當前元素小於堆頂元素，跳過該元素 否則替換堆頂元素爲當前元素，並重新調整堆 """ 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: 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) res = _.get_topk() print(sorted(res)) test()