LeeCode 1439 暴力 / 二分

题意

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

题解

暴力

枚举所有数组 $O(n^m)$，考虑到 $k$ 值较小，每一次枚举当前行，将子数组和大小压缩为 $min(k,tmp.size())$ ，只保留最小的 $k$ 个元素，暴力复杂度 $O(k\ast m\ast n)$。

class Solution {
public:
    int kthSmallest(vector<vector<int>>& mat, int k) {
        int m = mat.size(), n = mat[0].size();
        vector<int> res(mat[0].begin(), mat[0].end());
        for(int i = 1; i < m; i++){
            vector<int> tmp(res.size() * n);
            int cnt = 0;
            for(int j = 0; j < res.size(); j++){
                for(int k = 0; k < n; k++){
                   tmp[cnt++] = res[j] + mat[i][k];
                }
            }
            sort(tmp.begin(), tmp.end());
            res.assign(tmp.begin(), tmp.size() >= k ? tmp.begin() + k : tmp.end());
        }

        return res[k - 1];
    }
};

二分

二分答案，问题在于如何计算小于二分值的数据数量。题解中有个很妙的操作，预取每一行的最小值，这样保证 $dfs$ 时，所有行都已经有元素加入数组，此时只要增加同一行元素的差值，即可得到新的数组。

class Solution {
public:
    int m, n, K, cnt;
    void dfs(vector<vector<int>>& mat, int p, int sum, int mid){
        if(cnt >= K || p == m || sum >= mid) return;
        for(int i = 1; i < n; i++){
            if(sum + mat[p][i] - mat[p][0] < mid) ++cnt;
            dfs(mat, p + 1, sum + mat[p][i] - mat[p][0], mid);
        }
        dfs(mat, p + 1, sum, mid);
    }
    int kthSmallest(vector<vector<int>>& mat, int k) {
        K = k, m = mat.size(), n = mat[0].size();
        int lsum = 0, usum = 0;
        for(int i = 0; i < m; i++) lsum += mat[i][0], usum += mat[i][n - 1];
        int lb = lsum, ub = usum + 1;
        while(ub - lb > 1){
            int mid = (lb + ub) >> 1;
            cnt = 1;
            dfs(mat, 0, lsum, mid);
            if(cnt >= k) ub = mid;
            else lb = mid;
        }
        return lb;
    }
};