有序矩阵中的第k个最小数组和

Leetcode题目

给你一个 m * n 的矩阵 mat，以及一个整数 k ，矩阵中的每一行都以非递减的顺序排列。

你可以从每一行中选出 1 个元素形成一个数组。返回所有可能数组中的第 k 个 最小 数组和。

示例 1：

输入：mat = [[1,3,11],[2,4,6]], k = 5
输出：7
解释：从每一行中选出一个元素，前 k 个和最小的数组分别是：
[1,2], [1,4], [3,2], [3,4], [1,6]。其中第 5 个的和是 7 。
示例 2：

输入：mat = [[1,3,11],[2,4,6]], k = 9
输出：17
示例 3：

输入：mat = [[1,10,10],[1,4,5],[2,3,6]], k = 7
输出：9
解释：从每一行中选出一个元素，前 k 个和最小的数组分别是：
[1,1,2], [1,1,3], [1,4,2], [1,4,3], [1,1,6], [1,5,2], [1,5,3]。其中第 7 个的和是 9 。
示例 4：

输入：mat = [[1,1,10],[2,2,9]], k = 7
输出：12

这道题我本来想用回溯法列出所有可能，然后进行排序，果不其然超时了。

还有一种方法是这次重点要讲的，也就是使用堆结构+广度优先遍历的方法来求解。

具体思路：先把初始数组元素以及初始的数组和 加入小顶堆中，依次弹出堆中元素，弹出的元素代表了堆中存在的最小和的数组。当弹出到第k次时，即我们所要的答案。当每次弹出元素时，将该元素之和的所有可能数组放到堆中。

代码python版本

class Solution:
    def kthSmallest(self, mat: List[List[int]], k: int) -> int:
        m,n = len(mat),len(mat[0])
        pointer = [0]*m
        cur_sum = 0
        for i in range(m):
            cur_sum += mat[i][0]
        heap = []
        element = [cur_sum,tuple(pointer)]
        heapq.heappush(heap,element)
        repeat = set()
        repeat.add(tuple(pointer))
        idx = 0
        while idx<k:
            res,pointer = heapq.heappop(heap)
            for i,j in enumerate(pointer):
                if j<n-1:
                    new_pointer = list(pointer)
                    new_pointer[i] = j+1
                    new_pointer = tuple(new_pointer)
                    if new_pointer not in repeat:
                        cur_sum = 0
                        for x,y in enumerate(list(new_pointer)):
                            cur_sum+=mat[x][y]
                        element = [cur_sum,new_pointer]
                        heapq.heappush(heap,element)
                        repeat.add(new_pointer)
            idx+=1
        return res