1439. 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.

思路：优先队列，注意去重即可

class Solution(object):
    def kthSmallest(self, mat, k):
        """
        :type mat: List[List[int]]
        :type k: int
        :rtype: int
        """
        import heapq
        n,m=len(mat),len(mat[0])
        pq=[(sum(row[0] for row in mat),[0]*n)]
        vis=set([tuple(pq[0][1])])
        k-=1
        while k:
            v,s=heapq.heappop(pq)
            k -= 1
            # print(v, s)
            for si in range(len(s)):
                if s[si]+1>=m: continue
                ss = list(s)
                ss[si]+=1
                if tuple(ss) in vis: continue
                heapq.heappush(pq, (v+mat[si][ss[si]]-mat[si][ss[si]-1], ss))
                vis.add(tuple(ss))
        return pq[0][0]