Find the Kth Smallest Sum of a Matrix With Sorted Rows LeetCode solution in C++

1
class Solution {
2
 public:
3
  vector<int> kSmallestPairs(vector<int> &nums1, vector<int> &nums2, int k) {
4
    auto cmp = [&nums1, &nums2](pair<int, int> a, pair<int, int> b) {
5
      return nums1[a.first] + nums2[a.second] > nums1[b.first] + nums2[b.second];
6
    };
7
    int n = nums1.size();
8
    int m = nums2.size();
9
    vector<int> ans;
10
    if (n == 0 || m == 0) return ans;
11
    priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(cmp)> pq(cmp);
12
    pq.push({0, 0});
13
    while (k-- && !pq.empty()) {
14
      int i = pq.top().first;
15
      int j = pq.top().second;
16
      pq.pop();
17

18
      if (j + 1 < m) pq.push({i, j + 1});
19
      if (j == 0 && i + 1 < n) pq.push({i + 1, j});
20

21
      ans.push_back(nums1[i] + nums2[j]);
22
    }
23
    return ans;
24
  }
25
  int kthSmallest(vector<vector<int>> &mat, int k) {
26
    vector<int> nums1 = mat[0];
27
    int n = mat.size();
28
    for (int i = 1; i < n; ++i) {
29
      nums1 = kSmallestPairs(nums1, mat[i], k);
30
    }
31
    return nums1[k - 1];
32
  }
33
};