1550. Find The Kth Smallest Sum Of A Matrix With Sorted Rows¶

Difficulty: Hard

1550. Find the Kth Smallest Sum of a Matrix With Sorted Rows

Hard

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 k^th 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, n^m)
mat[i] is a non-decreasing array.

Solution¶

import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashSet;
import java.util.List;
import java.util.PriorityQueue;
import java.util.Set;

class Solution {
    public int kthSmallest(int[][] mat, int k) {
        int n = mat.length;
        List<Integer> curr = new ArrayList<>();
        for (int val : mat[0])
            curr.add(val);
        for (int i = 1; i < n; i++)
            curr = mergeKSmallest(curr, mat[i], k);
        return curr.get(k - 1);
    }

    private ArrayList<Integer> mergeKSmallest(List<Integer> a, int[] b, int k) {
        int m = a.size(), n = b.length;
        ArrayList<Integer> res = new ArrayList<>();

        PriorityQueue<int[]> pq = new PriorityQueue<>(Comparator.comparingInt(o -> o[0]));
        Set<String> visited = new HashSet<>();

        pq.offer(new int[] {a.get(0) + b[0], 0, 0});
        visited.add("0,0");

        while (!pq.isEmpty() && res.size() < k) {
            int[] top = pq.poll();
            int sum = top[0], i = top[1], j = top[2];
            res.add(sum);

            if (i + 1 < m && visited.add((i + 1) + "," + j)) {
                pq.offer(new int[] {a.get(i + 1) + b[j], i + 1, j});
            }
            if (j + 1 < n && visited.add(i + "," + (j + 1))) {
                pq.offer(new int[] {a.get(i) + b[j + 1], i, j + 1});
            }
        }

        return res;
    }
}

Complexity Analysis¶

Time Complexity: O(?)
Space Complexity: O(?)

Approach¶

Detailed explanation of the approach will be added here