记一道关于堆的算法题

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的时候同样取反即可。

后记：这其实是个面试题，好像还挺简单的？然而现场绕了半天，蠢哭……