LeetCode 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/description/

题目：

You are given an m x n matrix mat that has its rows sorted in non-decreasing order and an integer k.

You are allowed to choose exactly one 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.  

Constraints:

• m == mat.length
• n == mat.length[i]
• 1 <= m, n <= 40
• 1 <= mat[i][j] <= 5000
• 1 <= k <= min(200, nm)
• mat[i] is a non-decreasing array.

题解：

We know how to find the kth smallest sum for 2 arrays.

Now we have a matrix. We can start with 2, merge it into a new array and continue with rest rows of matrix.

Each time, we keep the k smallest sums into a new array and continue with the next row.

Time Complexity: O(m * (n + k) * logn)). m = mat.length. n = mat[0].length. 

Space: O(k + n).

AC Java:

class Solution {
    public int kthSmallest(int[][] mat, int k) {
        int[] row = mat[0];
        for(int r = 1; r < mat.length; r++){
            row = kSmall(row, mat[r], k);
        }

        return row[k - 1];
    }

    private int[] kSmall(int[] nums1, int[] nums2, int k){
        List<Integer> res = new ArrayList<>();
        int m = nums1.length;
        int n = nums2.length;
        PriorityQueue<int[]> minHeap = new PriorityQueue<>((a, b) -> Integer.compare(a[2], b[2]));
        for(int j = 0; j < n; j++){
            minHeap.add(new int[]{0, j, nums1[0] + nums2[j]});
        }

        while(k-- > 0 && !minHeap.isEmpty()){
            int[] cur = minHeap.poll();
            res.add(cur[2]);
            if(cur[0] == m - 1){
                continue;
            }

            minHeap.add(new int[]{cur[0] + 1, cur[1], nums1[cur[0] + 1] + nums2[cur[1]]});
        }

        int[] resArr = new int[res.size()];
        for(int i = 0; i < resArr.length; i++){
            resArr[i] = res.get(i);
        }

        return resArr;
    }
}

类似Find K Pairs with Smallest Sums.