Find K Pairs with Smallest Sums - Medium - L

2 minute read

Published:

Question

  • https://leetcode.com/problems/find-k-pairs-with-smallest-sums/

You are given two integer arrays nums1 and nums2 sorted in ascending order and an integer k.

Define a pair (u, v) which consists of one element from the first array and one element from the second array.

Return the k pairs (u1, v1), (u2, v2), …, (uk, vk) with the smallest sums.

Approach

  • Heaps

Solution

class Solution {
public:
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        vector<int> A = nums1;
        vector<int> B = nums2;
        int n = A.size();
        int K = k;
        vector<vector<int>> ans;
 
        // Min heap which contains tuple of the format
        // (sum, (i, j)) i and j are the indices
        // of the elements from array A
        // and array B which make up the sum.

        priority_queue<pair<int, pair<int, int> >,
                       vector<pair<int, pair<int, int> > >,
                       greater<pair<int, pair<int, int> > > > pq;

        // my_set is used to store the indices of
        // the  pair(i, j) we use my_set to make sure
        // the indices doe not repeat inside min heap.

        set<pair<int, int> > my_set;

        // initialize the heap with the minimum sum
        // combination i.e. (A[0] + B[0])
        // and also push indices (0,0) along
        // with sum.

        pq.push(make_pair(A[0] + B[0], make_pair(0, 0)));

        my_set.insert(make_pair(0, 0));

        // iterate upto K

        for (int count = 0; count < K; count++) {

            // tuple format (sum, i, j).
            pair<int, pair<int, int> > temp = pq.top();
            pq.pop();

            int i = temp.second.first;
            int j = temp.second.second;
            vector<int> t;
            t.push_back(A[i]);
            t.push_back(B[j]);
            ans.push_back(t);

            // check if i+1 is in the range of our first array A
            if (i + 1 < A.size()) {
                int sum = A[i + 1] + B[j];
                // insert (A[i + 1] + B[j], (i + 1, j))
                // into min heap.
                pair<int, int> temp1 = make_pair(i + 1, j);

                // insert only if the pair (i + 1, j) is
                // not already present inside the map i.e.
                // no repeating pair should be present inside
                // the heap.

                if (my_set.find(temp1) == my_set.end()) {
                    pq.push(make_pair(sum, temp1));
                    my_set.insert(temp1);
                }
            }
            // check if j+1 is in the range of our second array
            // B
            if (j + 1 < B.size()) {
                // insert (A[i] + B[j + 1], (i, j + 1))
                // into min heap.

                int sum = A[i] + B[j + 1];
                pair<int, int> temp1 = make_pair(i, j + 1);

                // insert only if the pair (i, j + 1)
                // is not present inside the heap.

                if (my_set.find(temp1) == my_set.end()) {
                    pq.push(make_pair(sum, temp1));
                    my_set.insert(temp1);
                }
            }
        }
        return ans;
    }
};