有序矩阵中第K小的元素（LeetCode）

news2024/9/22 15:35:37

题目

给你一个 n x n 矩阵 matrix ，其中每行和每列元素均按升序排序，找到矩阵中第 k 小的元素。
请注意，它是 排序后 的第 k 小元素，而不是第 k 个不同的元素。

你必须找到一个内存复杂度优于 $O(n^{2})$ 的解决方案。

解题

from queue import PriorityQueue


def find_kth_smallest(matrix, k):
    # 创建一个优先级队列，用于存储元素及其位置
    min_heap = PriorityQueue()

    # 将每一行的第一个元素及其位置 (值, 行, 列) 加入优先级队列
    for row in range(len(matrix)):
        min_heap.put((matrix[row][0], row, 0))

    kth_smallest = -1
    # 合并有序链表的逻辑，找到第 k 小的元素
    while not min_heap.empty() and k > 0:
        current_value, current_row, current_col = min_heap.get()
        # 更新第 k 小的元素
        kth_smallest = current_value
        k -= 1
        # 如果当前行还有更多的元素，则将下一个元素加入优先级队列
        if current_col + 1 < len(matrix[current_row]):
            next_value = matrix[current_row][current_col + 1]
            min_heap.put((next_value, current_row, current_col + 1))

    return kth_smallest


# 测试代码
matrix = [
    [1, 5, 9],
    [10, 11, 13],
    [12, 13, 15]
]
k = 8
print(f"这矩阵中第{k}小元素是: {find_kth_smallest(matrix, k)}")