题目要求返回所有可能数组中的第 k 个 最小 数组和。
根据题目给出的矩阵每行均为单调递增的提示,可以得出下一个最小数组和是由之前得到过的数组转移过来的。
因为k的范围不大,所以可以使用优先队列+bfs思想模拟这个过程。
备忘这里的优先队列重载方法
class Solution { public: int kthSmallest(vector<vector<int>>& mat, int k) { struct node { vector<int> pos; int tot; node(vector<int> tmp, int t) { pos = tmp, tot = t; } }; auto my_cmp = [](const node &a, const node &b) { return a.tot > b.tot; }; priority_queue<node, vector<node>, decltype(my_cmp)> pq(my_cmp); int M = mat.size(), N = mat[0].size(); vector<int> tmp(M); set<vector<int> > s; int sum = 0; for(int i = 0; i < M; i++) { tmp[i] = 0, sum += mat[i][0]; } pq.push({tmp, sum}); while(!pq.empty()) { tmp = pq.top().pos, sum = pq.top().tot; pq.pop(); k--; /*printf("%d ", sum); for(int i = 0; i < M; i++) printf("%d ", tmp[i]); puts("");*/ if(k == 0) break; for(int i = 0; i < M; i++) { if(tmp[i]+1 < N) { tmp[i]++; if(s.count(tmp) == 0) { s.insert(tmp); pq.push({tmp, sum+mat[i][tmp[i]]-mat[i][tmp[i]-1]}); } tmp[i]--; } } } return sum; } };