Python解决top-k问题

要从十亿行数据中找出最大的一百行数据,最常用的就是小根堆,而在python中小根堆的实现十分简单

class TopHeap(object):

    def __init__(self, k):
        self.k = k
        self.top_list = []

    def top(self, index):
        for i in index[:self.k]:
            heapq.heappush(self.top_list, i)
            heapq.heapify(self.top_list)

        for j in index[self.k:]:
            if j > self.top_list[0]:
                heapq.heapreplace(self.top_list, j)
                heapq.heapify(self.top_list)

        return self.top_list

定义一个TopHeap类,从带有众多数据的index数组中找出最大的前k个数据

index = np.random.randint(low=0, high=0x7FFFFFFF, size=1000000000)

使用numpy库来创建随机数组,如创建一个有十亿行数据的随机数组

    th = TopHeap(100)
    top_list = th.top(index)

然后通过TopHeap类找出这个数组中最大的前一百行数据

然而最终的运行时间与同样使用小根堆的c++或java的速度差了很多,所以如果还想进一步提高速度只能使用像CPython这样的方法来提高速度到c++的级别

猜你喜欢

转载自blog.csdn.net/weixin_55915644/article/details/125567542