python 实现 topk算法

python 实现 topk算法

这里的版本是针对的一个class的某一个成员变量进行的：
关于如何定义对象的比较方法，请参考往期文章：python定义对象的比较方法

class province_room_quality_data:
    def __init__(self, room, quality):
        self.room = room
        self.quality = quality
    def __lt__(self, other):
        return self.quality < other.quality

topk代码：

from typing import List
import heapq
# 对province_room_quality_data对象的TopK排序算法
def top_k(arr: List[province_room_quality_data], k: int) -> List[province_room_quality_data]:
    if k == 0:
        return list()
    hp = [item for item in arr[:k]]
    heapq.heapify(hp)
    for i in range(k, len(arr)):
        if hp[0].quality < arr[i].quality:
            heapq.heappop(hp)
            heapq.heappush(hp, arr[i])
    ans = [x for x in hp]
    return ans

最后可以得到基于quality的TopK大的对象数组，注意这里的对象数组并不是基于quality有序的，所以如果最后需要有序的话还得再执行一下sort

newList = top_k(list, 20)
newList.sort(reverse=True)