如何从排序矩阵中返回所有 k 个最小元素？

Question

我需要用伪代码解决这个问题：

给定一个 mXn 矩阵，其所有行和列均已排序。例如：

10 20 30 40 50
15 25 35 45 52
27 29 37 48 55
32 33 39 58 60

任务是用伪代码编写一个函数，以 k 作为参数，并返回所有 k 个最小元素。

例如，如果 k = 1 那么函数将返回 10

如果 k = 2 那么函数将返回 10 和 15（上面的矩阵只是一个例子）

k 是自然数并且 k <= m*n

所需时间复杂度为O(klogk)

我决定从矩阵元素构建一个最小堆，然后提取最小元素 k 次。

但是，构建堆的时间复杂度为 m*n 且为 k <= m*n my solution is not efficient enough.

解决这个问题更有效的算法是什么？

Answer 1

堆方法是一个好主意，但它没有利用每行和每列中的值已排序的事实。您可以从中受益。

考虑将行（不仅仅是单个值）放入堆中，以便堆中最多有 m 个条目。让堆由该堆中的行中最左边的值作为键。每次从该堆中提取一行时，从该行中提取最左边的值并将较短的行推回到堆上，现在由该行中新的“第一个”元素作为键。

这还不是最佳的，因为你的堆大小从一开始就是 m 。对于较小的 k，这不是最佳选择。

我们现在还可以使用已排序的列，并从具有一个条目的堆开始。每次从堆中提取一行（即堆上当前的“最后”行）时，也会将下一行添加到堆中。这样，堆每次迭代最多增长 1 个元素，因此您可以实现所需的复杂性。

此外，当您从堆中提取行时，您实际上不必更改行。实际上，您可以将堆条目设为（单元格值、单元格值坐标）的元组并使用它。

实施

这是 Python 中的实现：

from heapq import heappush, heappop

def get_min(matrix):
    m = len(matrix)
    if m == 0:
       return
    
    n = len(matrix[0])
    heap = [(matrix[0][0], 0, 0)]  # top-left element is minimum
    
    while heap:
        value, row, col = heappop(heap)
        yield value
        if col + 1 < n:
            heappush(heap, (matrix[row][col+1], row, col+1))
        # If we yielded a value from the first column, then also consider the value below it
        if col == 0 and row + 1 < m:
            heappush(heap, (matrix[row+1][0], row+1, 0))

这是一个生成器（“pythonic”），但如果您愿意，您也可以将

yield

语句替换为推送到结果列表，并在到达

k

时中断循环。生成器方法不需要这种中断，因为调用者有责任停止消耗生成的值。

示例运行可以这样编码：

matrix = [
    [10, 20, 30, 40, 50],
    [15, 25, 35, 45, 52],
    [27, 29, 37, 48, 55],
    [32, 33, 39, 58, 60],
]
k = 10

for i, value in enumerate(get_min(matrix)):
    if i >= k:
        break
    print(value)

这将输出：

10
15
20
25
27
29
30
32
33
35