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

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

這道題我本來想用回溯法列出所有可能，然後進行排序，果不其然超時了。

還有一種方法是這次重點要講的，也就是使用堆結構+廣度優先遍歷的方法來求解。

具體思路：先把初始陣列元素以及初始的陣列和 加入小頂堆中，依次彈出堆中元素，彈出的元素代表了堆中存在的最小和的陣列。當彈出到第k次時，即我們所要的答案。當每次彈出元素時，將該元素之和的所有可能陣列放到堆中。

程式碼python版本

 1 class Solution:
 2     def kthSmallest(self, mat: List[List[int]], k: int) -> int:
 3         m,n = len(mat),len(mat[0])
 4         pointer = [0]*m
 5         cur_sum = 0
 6         for i in range(m):
 7             cur_sum += mat[i][0]
 8         heap = []
 9         element = [cur_sum,tuple(pointer)]
 
10         heapq.heappush(heap,element)
11         repeat = set()
12         repeat.add(tuple(pointer))
13         idx = 0
14         while idx<k:
15             res,pointer = heapq.heappop(heap)
16             for i,j in enumerate(pointer):
17                 if j<n-1:
18                     new_pointer = list(pointer)
19                     new_pointer[i] = j+1
20                     new_pointer = tuple(new_pointer)
21                     if new_pointer not in repeat:
22                         cur_sum = 0
23                         for x,y in enumerate(list(new_pointer)):
24                             cur_sum+=mat[x][y]
25                         element = [cur_sum,new_pointer]
26                         heapq.heappush(heap,element)
27                         repeat.add(new_pointer)
28             idx+=1
29         return res