You are given an m * n
matrix, mat
, and an integer k
, which has its rows sorted in non-decreasing order.
You are allowed to choose exactly 1 element from each row to form an array. Return the Kth smallest array sum among all possible arrays.
Example 1:
Input: mat = [[1,3,11],[2,4,6]], k = 5 Output: 7 Explanation: Choosing one element from each row, the first k smallest sum are: [1,2], [1,4], [3,2], [3,4], [1,6]. Where the 5th sum is 7.
Example 2:
Input: mat = [[1,3,11],[2,4,6]], k = 9 Output: 17
Example 3:
Input: mat = [[1,10,10],[1,4,5],[2,3,6]], k = 7 Output: 9 Explanation: Choosing one element from each row, the first k smallest sum are: [1,1,2], [1,1,3], [1,4,2], [1,4,3], [1,1,6], [1,5,2], [1,5,3]. Where the 7th sum is 9.
Example 4:
Input: mat = [[1,1,10],[2,2,9]], k = 7 Output: 12
Constraints:
m == mat.length
n == mat.length[i]
1 <= m, n <= 40
1 <= k <= min(200, n ^ m)
1 <= mat[i][j] <= 5000
mat[i]
is a non decreasing array.
Related Topics:
Heap
Ugly answer wrote during the context.
// OJ: https://leetcode.com/problems/find-the-kth-smallest-sum-of-a-matrix-with-sorted-rows/
// Author: github.com/lzl124631x
typedef pair<int, string> Pair;
struct Cmp {
bool operator()(const Pair &a, const Pair &b) {
return a.first > b.first;
}
};
class Solution {
string hash(vector<int> &idx) {
string s;
for (int i = 0; i < idx.size(); ++i) {
auto h = to_string(idx[i]);
if (h.size() == 1) h = "0" + h;
s += h;
}
return s;
}
vector<int> dehash(string s, int M) {
vector<int> v(M);
for (int i = 0; i < M; ++i) v[i] = (s[2 * i] - '0') * 10 + s[2 * i + 1] - '0';
return v;
}
public:
int kthSmallest(vector<vector<int>>& A, int k) {
int M = A.size(), N = A[0].size(), sum = 0;
unordered_set<string> seen;
priority_queue<Pair, vector<Pair>, Cmp> q;
vector<int> idx(M);
auto idxHash = hash(idx);
seen.insert(idxHash);
for (int i = 0; i < M; ++i) sum += A[i][0];
q.emplace(sum, idxHash);
int x = 1;
while (k--) {
auto &p = q.top();
sum = p.first;
auto v = dehash(p.second, M);
for (int i = 0; i < M; ++i) {
if (v[i] + 1 >= N) continue;
int nsum = sum + A[i][v[i] + 1] - A[i][v[i]];
v[i]++;
auto h = hash(v);
if (!seen.count(h)) {
q.emplace(nsum, h);
seen.insert(h);
}
v[i]--;
}
q.pop();
}
return sum;
}
};