题目描述

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

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

样例

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

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

输入：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。 

输入：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] 是一个非递减数组。

算法

(枚举，优先队列) $O(km^2 \log (km))$

1. 维护一个优先队列，每个元素都是一个下标数组以及 sum，按照 sum 从小到大定义优先级。
2. 出队时，如果到了第 k 个元素，则返回当前元素的 sum。
3. 对于每个元素，都有 m 种方式进行拓展，即枚举每行都向后移动一个位置，判断新的元素是否已经访问过。如果没有访问过，则入队。

时间复杂度

• 每个元素有 $O(m)$ 个可拓展元素，队列和哈希表中元素个数最多为 $km$ 个，所以判断是否已经访问过需要 $O(m \log (km))$ 的时间，入队需要 $O(\log (km))$ 的时间。
• 循环这样的队列 $k$ 次，故总时间复杂度为 $O(km^2 \log (km))$。

空间复杂度

• 每个元素需要占用 $O(m)$ 的空间，故总共需要额外 $O(km^2)$ 的空间存储队列和哈希表。

C++ 代码

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

        #define PIV pair<int, vector<int>>

        priority_queue<PIV, vector<PIV>, greater<PIV>> q;
        set<vector<int>> vis;

        vector<int> idx(m, 0);
        int sum = 0;

        for (int i = 0; i < m; i++)
            sum += mat[i][0];

        q.push(make_pair(sum, idx));

        while (--k) {
            sum = q.top().first;
            idx = q.top().second;

            q.pop();

            for (int i = 0; i < m; i++) {
                int j = idx[i];
                if (j + 1 < n) {
                    idx[i]++;
                    if (vis.find(idx) == vis.end()) {
                        vis.insert(idx);
                        q.push(make_pair(sum - mat[i][j] + mat[i][j + 1], idx));
                    }
                    idx[i]--;
                }
            }
        }

        return q.top().first;
    }
};