python实现topk问题

方法1：
用python中的heapq实现

import heapq
import random


class TopKHeap(object):
    def __init__(self,k):
        self.data=[]
        self.k=k

    def push(self,num):

        if len(self.data)<self.k:
            heapq.heappush(self.data,num)
        else:
            top_min=self.data[0]
            if num>top_min:
                heapq.heapreplace(self.data,num)
    def topk(self):

        return list(reversed([heapq.heappop(self.data) for i in range(self.k)]))

if __name__ == '__main__':
    heap=TopKHeap(10)
    for i in range(100000):
        n=random.randint(0,100000)
        # print(n,'------------')
        heap.push(n)
    print(heap.topk())

方法2：
思路一样，就是自己写个最大堆，
步骤：

先用前k个数据建立好一个最小堆（k个结点）
然后对剩下的每一个数和根结点比较，比根结点大则替换根结点，替换之后此时就不是最小堆，所以需要再次建堆，使得变成最小堆；如过比根结点小，则继续遍历下一个数，直至遍历结束，
结束后，此时的堆就是这海量数据中最大的k个数，输出即可
注意：输出的数不是有序的，是最小堆的结果，要有序，得再进行排序
堆排序：https://blog.csdn.net/qq_41386300/article/details/89375074


import random


def heap(data,root):

    if 2*root+1<len(data):
        if 2*root+2<len(data) and data[2*root+2]<data[2*root+1]:
            k=2*root+2
        else:
            k=2*root+1
        if data[k]<data[root]:
            data[root],data[k]=data[k],data[root]
            heap(data,k)

def min_heap(data):
    for i in range(len(data)//2-1,-1,-1):
        heap(data,i)
    return data

def topk(data,k):

    data_k=min_heap(data[:k])
    for i in range(k,len(data)):
        if data[i]>data_k[0]:
            data_k[0]=data[i]
            data_k=min_heap(data_k)
    return data_k

if __name__ == '__main__':
    data=[random.randint(0,1000) for i in range(1000)]
    print(topk(data,10))

猜你喜欢