[LeetCode 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.

 

Extended problem of [LeetCode 373] Find K Pairs with Smallest Sums. Divide and conquer, then merge results by using the same approach in LeetCode 373. 

 

class Solution {
    public int kthSmallest(int[][] mat, int k) {
        int[] topK = compute(mat, k, 0, mat.length - 1);
        return topK[k - 1];
    }
    private int[] compute(int[][] mat, int k, int r1, int r2) {
        if(r1 == r2) {
            return mat[r1];
        }
        int mid = r1 + (r2 - r1) / 2;
        int[] top = compute(mat, k, r1, mid);
        int[] bottom = compute(mat, k, mid + 1, r2);
        int[] ans = new int[Math.min(k, top.length * bottom.length)];
        
        PriorityQueue<int[]> minPq = new PriorityQueue<>(Comparator.comparingInt(a -> top[a[0]] + bottom[a[1]]));
        for(int i = 0; i < Math.min(k, top.length); i++) {
            minPq.add(new int[]{i, 0});
        }
        int i = 0;
        while(i < ans.length) {
            int[] curr = minPq.poll();
            ans[i] = top[curr[0]] + bottom[curr[1]];
            i++;
            if(curr[1] < bottom.length - 1) {
                minPq.add(new int[]{curr[0], curr[1] + 1});
            }
        }
        return ans;
    }
}

 

 

Related Problems

[LeetCode 373] Find K Pairs with Smallest Sums