LeetCode 1439 - 有序矩阵中的第 k 个最小数组和（周赛）

题目描述

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

解法一：暴解（C++）

这里的排序工作我们借助了一下 multiset 来完成

class Solution {
public:
    int kthSmallest(vector<vector<int>>& mat, int k) {
        vector<int> res(mat[0]);
        int m = mat.size(), n = mat[0].size();
        for(int i=1;i<m;i++)
        {
            multiset<int> s;
            for(int x: res)
                for(int y: mat[i])
                    s.insert(x+y);
            res.assign(s.begin(), s.end());
            res.resize(min(k, (int)res.size()));
        }
        return res[k-1];
    }
};

解法二：二分查找（C++）

我们先确定左右边界，即最小和与最大和，然后二分得到 $mid$，每次判断和小于$mid$ 的组合有多少个，如果大于等于 $k$ 那么更新 $right$，否则更新$left$

这就是二分的内容了，但是这道题的关键点不仅仅在于二分，怎么找到可能的组合也是一个问题。

这里，我们采用了类似深搜的算法过程：先有一个初值 $sum$，然后从第 0 层一直往第 m-1 层搜索，$idx$ 表示当前在第几层，每次只改变自己所在的这一层的数，改变 $sum$， 如果小于 $mid$ 就往下一层搜，大于 $mid$ 就返回了

class Solution {
public:
    int kthSmallest(vector<vector<int>>& mat, int k) {
        int left = 0, right = 0;
        for(int i=0;i<mat.size();i++)
        {
            left += mat[i].front();
            right += mat[i].back();
        }

        int init = left;
        while(left<right)
        {
            int mid = left + (right-left)/2;
            int cnt = 1;
            dfs(mat, k, mid, 0, init, cnt);
            if(cnt<k) left = mid+1;
            else right = mid;
        }
        return right;
    }

    void dfs(vector<vector<int>>& mat, int k, int maxsum, int idx, int sum, int& cnt )
    {
        if(idx==mat.size()) return;
        if(sum>maxsum || cnt>k) return;
        
        dfs(mat, k, maxsum, idx+1, sum, cnt); // 这层就取首元素，故什么也不做，dfs向下一层

		// 不取首元素了，换个数取
        for(int j=1;j<mat[idx].size();j++)
        {
            int tmp = sum+mat[idx][j]-mat[idx][0];
            if(tmp>maxsum) break;
            cnt++;
            dfs(mat, k, maxsum, idx+1, tmp, cnt);
        }
    }
};