Find the Kth Smallest Sum of a Matrix With Sorted Rows LeetCode solution in Python

1
class Solution:
2
    def kthSmallest(self, mat: List[List[int]], k: int) -> int:
3
        def kSmallestPairs(
4
            nums1: List[int], nums2: List[int], k: int
5
        ) -> List[List[int]]:
6
            h = [(nums1[0] + nums2[0], 0, 0)]
7
            visited = set()
8
            res = []
9
            while h and k > 0:
10
                e, i, j = heappop(h)
11
                if (i, j) in visited:
12
                    continue
13
                res.append(e)
14
                visited.add((i, j))
15
                if j + 1 < len(nums2):
16
                    heappush(h, (nums1[i] + nums2[j + 1], i, j + 1))
17
                if i + 1 < len(nums1):
18
                    heappush(h, (nums1[i + 1] + nums2[j], i + 1, j))
19
                k -= 1
20
            return res
21

22
        res = mat[0]
23
        for i in range(1, len(mat)):
24
            res = kSmallestPairs(res, mat[i], k)
25
        return res[-1]