Skip to content

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

SourceWeekly Contest 187 Q4DifficultyHardRating2133

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.

Solutions

Solution 1

Thinking

One value per row yields \(n^m\) sums, but \(k\le 200\). We only need the \(k\) smallest partial sums. Rows are sorted, so combining with the next row uses at most its first \(k\) entries.

Let \(pre\) be the \(k\) smallest sums so far. Cartesian-add the current row, sort, and keep \(k\) values. After the last row, \(pre[k-1]\) is the answer.

1
2
3
4
5
6
class Solution:
    def kthSmallest(self, mat: List[List[int]], k: int) -> int:
        pre = [0]
        for cur in mat:
            pre = sorted(a + b for a in pre for b in cur[:k])[:k]
        return pre[-1]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
class Solution {
    public int kthSmallest(int[][] mat, int k) {
        int m = mat.length, n = mat[0].length;
        List<Integer> pre = new ArrayList<>(k);
        List<Integer> cur = new ArrayList<>(n * k);
        pre.add(0);
        for (int[] row : mat) {
            cur.clear();
            for (int a : pre) {
                for (int b : row) {
                    cur.add(a + b);
                }
            }
            Collections.sort(cur);
            pre.clear();
            for (int i = 0; i < Math.min(k, cur.size()); ++i) {
                pre.add(cur.get(i));
            }
        }
        return pre.get(k - 1);
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
class Solution {
public:
    int kthSmallest(vector<vector<int>>& mat, int k) {
        int pre[k];
        int cur[mat[0].size() * k];
        memset(pre, 0, sizeof pre);
        int size = 1;
        for (auto& row : mat) {
            int i = 0;
            for (int j = 0; j < size; ++j) {
                for (int& v : row) {
                    cur[i++] = pre[j] + v;
                }
            }
            sort(cur, cur + i);
            size = min(i, k);
            for (int j = 0; j < size; ++j) {
                pre[j] = cur[j];
            }
        }
        return pre[k - 1];
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
func kthSmallest(mat [][]int, k int) int {
    pre := []int{0}
    for _, row := range mat {
        cur := []int{}
        for _, a := range pre {
            for _, b := range row {
                cur = append(cur, a+b)
            }
        }
        sort.Ints(cur)
        pre = cur[:min(k, len(cur))]
    }
    return pre[k-1]
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
function kthSmallest(mat: number[][], k: number): number {
    let pre: number[] = [0];
    for (const cur of mat) {
        const next: number[] = [];
        for (const a of pre) {
            for (const b of cur) {
                next.push(a + b);
            }
        }
        pre = next.sort((a, b) => a - b).slice(0, k);
    }
    return pre[k - 1];
}

Comments