Leetcode題解 Python：Find the Kth Smallest Sum of a Matrix With Sorted Rows

給一個 m*n 的矩陣，跟整數 k，而且每行已經排列。從每行取一個出來作總和，求第 k 小的總和是多少。

這題跟 373. 是很類似的，不過這題多了很多行數列，但並不是什麼難點。

若採用全部枚舉的話，最多40**40個，這樣子會超時。

若分批進行，一樣從全部的索引 0 開始，每輪後索引總和+1，繼續枚舉，因為目標大小 k 最高200個，所以這方法是可行的？

只要每輪枚舉結束，枚舉的總數超過 k 就可以找到目標值？但並不是這麼一回事。

的確一開是從最小值開始枚舉，會得到各索引+1的組合。

若只是單純以索引+1為枚舉的方式，這不能保證一定能包含完整的最小總和順序。

例如：mat = [[1,2,3,4,5,6],[1,1,1,1,1,1]], k = 3

然而下一輪要枚舉的，是以什麼當基準呢？

如果一開始是最小值，那麼，第二小的值會在枚舉中被找到。

若換種角度，都以最小值的索引組合為基準，而第三小的值會在第二小的枚舉中被找到，第四小的也是同理。

因此我們需要參考最小值的索引組合做枚舉。使用 heap 的話，就能保持最小值總在第一個，也方便進行後面的枚舉。

如此連鎖，第二找第三，第三找第四，……，就能找到第 k 的值。

Python

class SolutionPass:
    def kthSmallest(self, mat, k: int) -> int:
        import heapq
        result = []
        n, m = len(mat), len(mat[0])
        heap = [(sum(each[0] for each in mat) , *[0 for _ in range(n)])]
        memo = set()             
        while k > 0 and heap:
            sm, *idxs = heapq.heappop(heap)
            result.append(sm)
            k -= 1
            for i in range(n):
                idxs[i] += 1
                if idxs[i] < m and tuple(idxs) not in memo:       
                    memo.add(tuple(idxs))
                    heapq.heappush(heap, [sum([mat[j][idxs[j]] for j in range(n)]), *idxs])
                idxs[i] -= 1
        return result[-1]