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/

将matrix降维到373 两个num array取 k smallest pair

373 是在Priority Queue里面做一个BFS,1439是将BFS的结果与每一个新的row融合。

class Solution:
    def kthSmallest(self, mat: List[List[int]], k: int) -> int:
        def kSmallestPairs(nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
            queue = []
            res = []
            visited = dict()
            if len(nums1) == 0 or len(nums2) == 0:
                return res
            heapq.heappush(queue,(nums1[0]+nums2[0],0,0))
            visited[(0,0)] = 1
            while len(res) < k and queue:
                curr = heapq.heappop(queue)
                currSum,i,j = curr
                res.append(nums1[i]+nums2[j])
                if i+1 < len(nums1) and (i+1,j) not in visited:
                    heapq.heappush(queue,(nums1[i+1]+nums2[j],i+1,j))
                    visited[(i+1,j)] = 1
                if j+1 < len(nums2) and (i,j+1) not in visited:
                    heapq.heappush(queue,(nums1[i]+nums2[j+1],i,j+1))
                    visited[(i,j+1)] = 1
            return res
        curr = mat[0]
        for row in mat[1:]:
            curr = kSmallestPairs(curr,row,k)
        return curr[-1]

Last updated

Was this helpful?