【leetcode C++】1439. 有序矩阵中的第 k 个最小数组和 1439. Find the Kth Smallest Sum of a Matrix With Sorted Rows

1439. 有序矩阵中的第 k 个最小数组和
1439. Find the Kth Smallest Sum of a Matrix With Sorted Rows

class Solution {
public:
    int kthSmallest(vector<vector<int>>& mat, int k) {
        priority_queue<pair<int, vector<int> > > PQ;  // 优先队列，用来存放sum及其路径
        set<vector<int> > S;  // 路径的集合，防止重复
        int sum = 0, result = 0;
        vector<int> path;  // 当前路径
        for(int ii = 0; ii < mat.size(); ii++) {  // 初始化
            sum += mat[ii][0];
            path.push_back(0);
        }
        PQ.push({-sum, path});  // 对sum取相反数，用大顶堆来模拟小顶堆
        for(int cnt = 0; cnt < k; cnt++) {
            auto A = PQ.top();  // 取出第 cnt 个最小值
            PQ.pop();
            result = -A.first;
            path = A.second;
            for(int ii = 0; ii < path.size(); ii++) {  // 放入下一次所有可能的最小值到优先队列中
                if(path[ii] + 1 >= mat[ii].size()) continue;  // 这里注意，路径下标不能超过当前数组大小
                path[ii]++;
                if(S.find(path) == S.end()) {
                    sum = 0;
                    for(int jj = 0; jj < path.size(); jj++) {
                        sum += mat[jj][path[jj]];
                    }
                    PQ.push({-sum, path});
                    S.insert(path);
                }
                path[ii]--;
            }
        }
        return result;
    }
};