跳转至

1439. 有序矩阵中的第 k 个最小数组和

来源第 187 场周赛 Q4难度困难分数2133

题目描述

给你一个 m * n 的矩阵 mat，以及一个整数 k ，矩阵中的每一行都以非递减的顺序排列。

你可以从每一行中选出 1 个元素形成一个数组。返回所有可能数组中的第 k 个 最小 数组和。

 

示例 1：

输入：mat = [[1,3,11],[2,4,6]], k = 5
输出：7
解释：从每一行中选出一个元素，前 k 个和最小的数组分别是：
[1,2], [1,4], [3,2], [3,4], [1,6]。其中第 5 个的和是 7 。

示例 2：

输入：mat = [[1,3,11],[2,4,6]], k = 9
输出：17

示例 3：

输入：mat = [[1,10,10],[1,4,5],[2,3,6]], k = 7
输出：9
解释：从每一行中选出一个元素，前 k 个和最小的数组分别是：
[1,1,2], [1,1,3], [1,4,2], [1,4,3], [1,1,6], [1,5,2], [1,5,3]。其中第 7 个的和是 9 。

示例 4：

输入：mat = [[1,1,10],[2,2,9]], k = 7
输出：12

 

提示：

  • m == mat.length
  • n == mat.length[i]
  • 1 <= m, n <= 40
  • 1 <= k <= min(200, n ^ m)
  • 1 <= mat[i][j] <= 5000
  • mat[i] 是一个非递减数组

解法

方法一：逐行遍历 + 排序

思考

每行选一个数求和，\(n^m\) 种组合而 \(k\le 200\)。只需保留前 \(k\) 小的部分和。行已有序，与下一行相加时只需用到该行前 \(k\) 个元素。

设 \(pre\) 为已处理行的前 \(k\) 小和，与当前行做笛卡尔加后排序截断前 \(k\) 项。处理完所有行后 \(pre[k-1]\) 即为答案。

根据题目描述，我们需要找出前 \(m\) 行的所有可能数组中的第 \(k\) 个最小数组和。

如果我们能够找出前 \(m - 1\) 行的所有可能数组中的前 \(k\) 个最小数组和，那么我们可以将第 \(m\) 行的每个元素与前 \(m - 1\) 行的前 \(k\) 个最小数组和相加，将得到的所有结果排序后，取前 \(k\) 个最小值，即为前 \(m\) 行的所有可能数组中的前 \(k\) 个最小值。

因此，我们可以定义一个数组 \(pre\)，用于存储此前遍历到的行的前 \(k\) 个最小数组和，初始时 \(pre\) 只有一个元素 \(0\)。

然后，我们遍历 \(mat\) 的每一行 \(cur\)，将 \(cur\) 中的每个元素与 \(pre\) 中的每个元素相加，将得到的所有结果排序后，取前 \(k\) 个最小值作为新的 \(pre\)。继续遍历下一行，直到遍历完所有行。

最后返回 \(pre\) 中的第 \(k\) 个数（下标 \(k-1\)）即可。

时间复杂度 \(O(m \times n \times k \times \log (n \times k))\)，空间复杂度 \(O(n \times k)\)。其中 \(m\) 和 \(n\) 分别是矩阵的行数和列数。

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];
}

评论