LC-1439. 有序矩阵中的第 k 个最小数组和（二分答案、多路归并）

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

难度困难120

给你一个 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] 是一个非递减数组

二分查找解决第k小问题

https://leetcode.cn/problems/find-the-kth-smallest-sum-of-a-matrix-with-sorted-rows/solution/by-lfool-w2ul/

719. 找出第 K 小的数对距离

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

我们先明确一下原问题：「返回所有可能数组中的第 k 个最小数组和」

对于一个给定的数组和 sum，如果该数组和小了，那我们就「向右收缩」；如果该数组和大了，那我们就「向左收缩」；如果该数组和满足要求，那我们能不能找一个更小的且满足要求的数组和，所以需要继续「向左收缩」(是不是很像「寻找最左相等元素」？自信点，就是「寻找最左相等元素」)

所以「搜索对象」是什么？？很明显就是「数组和」嘛！！

那「搜索范围」又是什么呢？？

• 数组和最小值可以到达多少？肯定是该矩形的第一列组成的组数和嘛 (矩形每行递增)
• 数组和最大值可以到达多少？肯定是该矩形的最后一列组成的组数和嘛 (矩形每行递增)

明确了「搜索对象」和「搜索范围」，我们还需要搞清楚怎么确定数组和小了还是大了，肯定要有一个参考对象才能确定近还是远嘛

很聪明，这个参考对象就是题目给的「第 k 最小」。对于数组和 sum，小于等于该数组和的数量为 n，如果 n < k 说明数组和小了；如果 n > k 说明数组和大了

class Solution {
    // https://leetcode.cn/problems/find-the-kth-smallest-sum-of-a-matrix-with-sorted-rows/solution/by-lfool-w2ul/
    int m, n, k;
    int[][] mat;
    int cnt = 0; // 计算小于等于当前数组和的数量
    public int kthSmallest(int[][] mat, int k) {
        this.k = k;
        this.mat = mat;
        m = mat.length;
        n = mat[0].length;
        // 搜索范围
        int left = 0, right = 0;
        for (int i = 0; i < m; i++) {
            left += mat[i][0];
            right += mat[i][n - 1];
        }
        // 把最小值设为初始值
        int init = left;
        while (left <= right) {
            int mid = left - (left - right) / 2;
            // 初始值也算一个可行解
            cnt = 1;
            dfs(0, init, mid);
            // 对应数组和大了，向左收缩
            if (cnt >= k) right = mid - 1;
            // 对应数组和小了，向右收缩
            else left = mid + 1;
        }
        return left;
    }

    // DFS 计算 小于等于 target 的数量
    private void dfs(int row, int sum, int target) {
        // 特殊情况，直接返回
        // sum > target：当前数组和大于 target
        // cnt > k：当前小于等于 target 的数量大于 k
        // row >= m：已经到达最后一行 (结束条件)
        if (sum > target || cnt > k || row >= m) return;
        // 不做交换
        dfs(row + 1, sum, target);
        // 分别与 [1, n-1] 交换
        for (int i = 1; i < n; i++) {
            // 更新数组和：减去「第一个元素」，加上「要交换的元素」
            int newSum = sum - mat[row][0] + mat[row][i];
            // 交换后的数组和大于 target，直接跳出循环
            // 原因：由于每行元素递增，所以当前元素大了，该行后面的元素肯定也大了
            if (newSum > target) break;
            // 满足要求，cnt + 1
            cnt++;
            // 搜索下一行
            dfs(row + 1, newSum, target);
        }
    }
}

多路归并

题解：https://leetcode.cn/problems/find-the-kth-smallest-sum-of-a-matrix-with-sorted-rows/solution/by-lfool-z9n4/

class Solution {
    // 对于每次弹出队顶元素后，需要把 n 个元素入队
    public int kthSmallest(int[][] mat, int k) {
        int m = mat.length, n = mat[0].length;
        Set<String> set = new HashSet<>();
        // 构造一个最小堆
        PriorityQueue<int[]> q = new PriorityQueue<>((a, b) -> a[m] - b[m]);
        // init[0..m-1] 存储矩阵的一列元素的下标
        // init[m]存储该列的和
        int[] init = new int[m+1];
        for(int i = 0; i < m; i++){
            init[m] += mat[i][0];
            init[i] = 0;
        }
        q.offer(init);
        set.add(Arrays.toString(init));
        while(k-- > 0){
            int[] cur = q.poll();
            if(k == 0) return cur[m];
            // 构造处需要加入堆的元素，一共m个
            for(int i = 0; i < m; i++){
                int[] tmp = (int[])Arrays.copyOf(cur, m+1);
                if(tmp[i] + 1 >= n)
                    continue; // 下标越界了
                tmp[m] += mat[i][tmp[i] + 1] - mat[i][tmp[i]]; // 下标i往前进，tmp[m]处增加上变化量
                tmp[i] += 1;
                String s = Arrays.toString(tmp);
                if(set.contains(s)) continue;
                q.offer(tmp);
                set.add(s);
            }
        }
        return -1;
    }
}

多路归并练习题

264. 丑数 II

313. 超级丑数

373. 查找和最小的 K 对数字

786. 第 K 个最小的素数分数

1508. 子数组和排序后的区间和

719. 找出第 K 小的数对距离

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