1 题目描述
有20个数组,每个数组中有500个数,降序排列,输出这20个数组中最大的500个数。
2 解法
我开始想的是归并的方法。每个数组用一个指针指向第一个元素,然后比较得到最大值,指向最大值的指针后移,一共循环500次。但是这样的时间复杂度是O(n*k),可以用堆将其降为O(nlogk)。(假如数组有k个,每个数组中有n个元素)
具体做法是先取每个数组的首元素建立一个大小为k的大顶堆,pop出最大值保存到res数组,然后push下一个数进去。这里的下一个数就是最大值在数组中的后面一个数。因此我们需要在建堆的同时记录数的下标。为了方便,堆的元素就用元组表示。
代码如下:
def getTopk(matrix):
import heapq
res = []
m, n = len(matrix), len(matrix[0])
heap = [(-matrix[i][0], i, 0) for i in range(m)]
heapq.heapify(heap)
for i in range(n):
value, x, y = heapq.heappop(heap)
res.append(-value)
heapq.heappush(heap, (-matrix[x][y + 1], x, y + 1))
return res
注意,python中heapq默认建的是最小堆。所以push的时候需要对数取反,pop的时候同样取反即可。
后记:这其实是个面试题,好像还挺简单的?然而现场绕了半天,蠢哭……