You are given an m * n matrix, mat, and an integer k, which has its rows sorted in non-decreasing order.
You are allowed to choose exactly 1 element from each row to form an array. Return the Kth smallest array sum among all possible arrays.
Example 1:
Input: mat = [[1,3,11],[2,4,6]], k = 5
Output: 7
Explanation: Choosing one element from each row, the first k smallest sum are:
[1,2], [1,4], [3,2], [3,4], [1,6]. Where the 5th sum is 7.
Example 2:
Input: mat = [[1,3,11],[2,4,6]], k = 9
Output: 17
Example 3:
Input: mat = [[1,10,10],[1,4,5],[2,3,6]], k = 7
Output: 9
Explanation: Choosing one element from each row, the first k smallest sum are:
[1,1,2], [1,1,3], [1,4,2], [1,4,3], [1,1,6], [1,5,2], [1,5,3]. Where the 7th sum is 9.
Example 4:
Input: mat = [[1,1,10],[2,2,9]], k = 7
Output: 12
Constraints:
m == mat.length
n == mat.length[i]
1 <= m, n <= 40
1 <= k <= min(200, n ^ m)
1 <= mat[i][j] <= 5000
mat[i] is a non decreasing array.
在每行都排好序的M*N矩阵中每行抽一个元素组成新的数组，求所有可能的数组中第K小的数组和。是373. Find K Pairs with Smallest Sums的升级版，将两个数组增加到M个数组，思路是一致的，在两个数组中找到前K小的『pair和』cand后，与新数组的第K小只会是cand与新数组再求Find K Pairs with Smallest Sums，因为cand外的即使与nums3[0]相加也会超出前K小，此时得到的是三个数组中前K小的3-tuples；依次类推直到所有M个数组都计算一遍。
from heapq import heappop, heappush
class Solution:
def kthSmallest(self, mat: List[List[int]], k: int) -> int:
def kth(nums1, nums2, K):
k = K
res, q = [], []
for i in range(k):
if i == len(nums1):
break
heappush(q, (nums1[i] + nums2[0], nums1[i], 0))
while k > 0 and q:
cur = heappop(q)
res.append(cur[0])
k -= 1
if cur[2] == len(nums2) - 1:
continue
heappush(q, (cur[1] + nums2[cur[2] + 1], cur[1], cur[2] + 1))
return res
cand = mat[0]
for i in range(1, len(mat)):
cand = kth(cand, mat[i], k)
return cand[k - 1]