373. Find K Pairs with Smallest Sums - cocoder39/coco39_LC GitHub Wiki

373. Find K Pairs with Smallest Sums

cited from leetcode discussion

Suppose you are at pair 0,0 (index 0 and index 0, not value), which is the initial minimum. Then you know the only two possible followers (immediate larger ones) are 0,1 and 1,0. Any other indices, say 1,1, 1,2, 3,3 have to be larger. So every time you get a current minimum i,j , you want to push i+1,j and i,j+1 into heap. You don't need to worry about others yet.

The problem here is, if you are at 2,3, you will push 3,3 and 2,4; then later, you are at 3,2. Then you will push 3,3, 4,2. so you pushed 3,3 twice.

But it is easy to realize, if current minimum is 2,3 if you haven't seen 3,2 yet (meaning 3,2 is still in the queue, never been a current minimum yet), you don't need to worry about 3,3 at this moment, because 3,2 is guaranteed to be no less than 3,3. So you can wait till you see 3,2.

when current minimum is (i, j), the next minimum can be an element within pq or (i + 1, j), or (i, j + 1). And (i + 1, j) and (i, j + 1) might be in pq as well. Hence, pushing both (i, j + 1) and (i + 1, j) into pq can produce duplicate. we can avoid duplicate if we follow the style that only push (i + 1, j) into pq because (i, j + 1) would follow (i - 1, j + 1). if (i - 1, j + 1) has not been processed yet, then there is no need to push (i, j + 1); otherwise since (i - 1, j + 1) has been processed, then (i, j + 1) is already in pq

This basically means every time you see i,j you just need to push i+1, j into queue. i, j+1 can be pushed into queue later when you see i - 1, j + 1. The only exception to this is, when i == 0, there is no i-1, j+1 anymore, so you want to push both i+1, j and i, j+1 when i == 0.

we can follow another style that only (i, j + 1) follows (i, j)

tip: comparator in pq

auto cmp = [&nums1, &nums2](pair<int, int> p1, pair<int, int> p2) {
            return nums1[p1.first] + nums2[p1.second] > nums1[p2.first] + nums2[p2.second];
        };
        priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(cmp)> pq(cmp);

each time increase at most 2 elements, pop out 1 element. the size is no more than k. time is O(k log k), memory is O(k)

vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        vector<pair<int, int>> res;
        if (nums1.empty() || nums2.empty() || k <= 0) {
            return res;
        }
        
        auto cmp = [&nums1, &nums2](pair<int, int> p1, pair<int, int> p2) {
            return nums1[p1.first] + nums2[p1.second] > nums1[p2.first] + nums2[p2.second];
        };
        priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(cmp)> pq(cmp);
        
        pq.push(make_pair(0, 0));
        while (! pq.empty() && res.size() < k) {
            auto p = pq.top();
            pq.pop();
            res.push_back(make_pair(nums1[p.first], nums2[p.second]));
            
            if (p.first + 1 < nums1.size()) {
                pq.push(make_pair(p.first + 1, p.second));
            }
            if (p.first == 0 && p.second + 1 < nums2.size()) {
                pq.push(make_pair(p.first, p.second + 1));
            }
        }
        return res;
    }