二分 + dfs剪枝

时间复杂度

二分部分是log(m * 5000) 因为1 <= mat[i][j] <= 5000
dfs部分就有点 玄学

class Solution {
public:
    int m, n;
    int kthSmallest(vector<vector<int>>& mat, int k) {
        m = mat.size(), n = mat[0].size();

        int left = 0, right = 0;
        // 二分答案, left必定是最左一列, right必定是最右一列
        for(int i = 0; i < m; i ++) left += mat[i][0], right += mat[i][n - 1];

        while (left < right) {  
            // 模板一      
            int mid = left + right >> 1;
            int cnt = counter(mat, mid, 0, 0, k);
            if(cnt >= k) right = mid;
            else left = mid + 1;
        }
        return left;
    }

    // 搜索, 
    int counter(vector<vector<int>>& mat, int targetSum, int row, int sum, int k) {
        if (sum > targetSum) return 0;
        if (row == m) return 1;
        int ans = 0;
        for (int c = 0; c < n; c ++) {
            // ans, 这里是统计满足 sum <= targetSum 的个数, 
            // 上一次搜索的结果和下一个搜索的结果是不相交的, 故需要减去ans, 之后的搜索只需要搜到k-ans个即可, 多了不行
            int cnt = counter(mat, targetSum, row + 1, sum + mat[row][c], k - ans);
            // 若当前这一列都不满足, 则必然在之前的列中, 而不是之后的列中
            if (cnt == 0) break;
            ans += cnt;
            if (ans > k) break; // prune when ans > k
        }
        return ans;
    }
};