1439. Find the Kth Smallest Sum of a Matrix With Sorted Rows

https://leetcode.com/problems/find-the-kth-smallest-sum-of-a-matrix-with-sorted-rows/

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]