如何使用 Python 堆实现有序矩阵中第 K 小的元素？

2023-04-11 00:00:00 有序如何使用矩阵

利用 Python 的堆实现有序矩阵中第 K 小的元素可以通过以下步骤实现：

1.将有序矩阵的第一行元素添加到堆中；

2.从堆中取出最小元素，并将其所在列的下一个元素添加到堆中，直到取出了第 k 小的元素。

下面是代码演示：

import heapq

def kthSmallest(matrix, k):
    heap = [(matrix[0][j], 0, j) for j in range(len(matrix))]
    heapq.heapify(heap)
    count = 0
    while heap:
        val, row, col = heapq.heappop(heap)
        count += 1
        if count == k:
            return val
        if row < len(matrix) - 1:
            heapq.heappush(heap, (matrix[row + 1][col], row + 1, col))

假设有序矩阵如下：

[
 [ 1,  5,  9],
 [10, 11, 13],
 [12, 13, 15]
]

如果要找到矩阵中第 8 小的元素，可以用以下代码调用函数：

matrix = [
    [1, 5, 9],
    [10, 11, 13],
    [12, 13, 15]
]
k = 8
result = kthSmallest(matrix, k)
print(result)  # 输出 13

经过堆排序后，第 8 小的元素是 13。

相关文章