1439. Find the Kth Smallest Sum of a Matrix With Sorted Rows

O(mklog(min(n, k))) time O(min(n, k)) space
373. Find K Pairs with Smallest Sums的followup
不要想复杂了！！！直接做m - 1遍就行了
利用这种原理还可以做并行的两两merge

class Solution {
    vector<int> kSmallestPairs(const vector<int>& nums1, const vector<int>& nums2, int k) {
        if (nums1.empty() || nums2.empty() || k == 0) return {};
        using pii = pair<int, int>;
        auto cmp = [&nums1, &nums2](const pii &a, const pii &b) { return nums1[a.first] + nums2[a.second] > nums1[b.first] + nums2[b.second]; };
        priority_queue<pii, vector<pii>, decltype(cmp)> q(cmp);
        int m = nums1.size(), n = nums2.size();
        for (int j = 0, sz = min(k, n); j < sz; ++j) {
            q.emplace(0, j);
        }
        vector<int> res;
        while (k-- > 0 && !q.empty()) {
            auto [i, j] = q.top(); q.pop();
            res.push_back(nums1[i] + nums2[j]);
            if (i + 1 < m) {
                q.emplace(i + 1, j);
            }
        }
        return res;
    }
public:
    int kthSmallest(vector<vector<int>>& mat, int k) {
        int m = mat.size(), n = min((int)mat[0].size(), k);
        vector<int> v(begin(mat[0]), begin(mat[0]) + n);
        for (int r = 1; r < m; ++r) {
            v = kSmallestPairs(v, mat[r], k);
        }
        return v.back();
    }
};

并行merge版
跟23. Merge k Sorted Lists思路一样

class Solution {
    vector<int> kSmallestPairs(const vector<int>& nums1, const vector<int>& nums2, int k) {
        if (k == 0) return {};
        if (nums1.empty()) return nums2;
        if (nums2.empty()) return nums1;
        using pii = pair<int, int>;
        auto cmp = [&nums1, &nums2](const pii &a, const pii &b) { return nums1[a.first] + nums2[a.second] > nums1[b.first] + nums2[b.second]; };
        priority_queue<pii, vector<pii>, decltype(cmp)> q(cmp);
        int m = nums1.size(), n = nums2.size();
        for (int j = 0, sz = min(k, n); j < sz; ++j) {
            q.emplace(0, j);
        }
        vector<int> res;
        while (k-- > 0 && !q.empty()) {
            auto [i, j] = q.top(); q.pop();
            res.push_back(nums1[i] + nums2[j]);
            if (i + 1 < m) {
                q.emplace(i + 1, j);
            }
        }
        return res;
    }
public:
    int kthSmallest(vector<vector<int>>& mat, int k) {
        int m = size(mat);
        for (int step = 1; step < m; step <<= 1) {
            for (int i = 0; i + step < m; i += (step << 1)) {
                mat[i] = kSmallestPairs(mat[i], mat[i + step], k);
            }
        }
        return mat[0].back();
    }
};