algoadvance

Leetcode 373. Find K Pairs with Smallest Sums

Problem Statement

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 (u, v) with the smallest sums.

Clarifying Questions

  1. Can nums1 and nums2 contain negative numbers?
    • Yes, both arrays can contain negative, zero, or positive integers.
  2. What should be done if k is greater than the product of the lengths of the two arrays?
    • In such a case, we would return all possible pairs.
  3. Are the elements in nums1 and nums2 guaranteed to be sorted?
    • Yes, both nums1 and nums2 are sorted in ascending order.

Strategy

  1. Use a min-heap to keep track of the smallest sums.
  2. Initialize the heap with the smallest possible pairs (nums1[i], nums2[0]) for all valid i.
  3. Extract the smallest pair from the heap and add its next potential pair into the heap to maintain the k smallest pairs.
  4. Continue this process until we have k pairs or exhaust all possible pairs.

Code Implementation

#include <vector>
#include <queue>
#include <utility>

using namespace std;

class Solution {
public:
    vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        vector<pair<int, int>> result;
        if (nums1.empty() || nums2.empty() || k == 0)
            return result;

        auto compare = [&](pair<int, int>& a, pair<int, int>& b) {
            return nums1[a.first] + nums2[a.second] > nums1[b.first] + nums2[b.second];
        };
        
        priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(compare)> min_heap(compare);

        // Initialize the heap with the first column's pairs.
        for (int i = 0; i < nums1.size() && i < k; ++i) {
            min_heap.push({i, 0});
        }

        // Extract the smallest pairs.
        while (k-- > 0 && !min_heap.empty()) {
            auto idx_pair = min_heap.top();
            min_heap.pop();
            int i = idx_pair.first;
            int j = idx_pair.second;
            result.push_back({nums1[i], nums2[j]});

            if (j + 1 < nums2.size()) {
                min_heap.push({i, j + 1});
            }
        }

        return result;
    }
};

Time Complexity

Thus, the overall time complexity is O(min(k, len(nums1) * len(nums2)) * log k).

This ensures that our approach efficiently finds the k smallest pairs.

Cut your prep time in half and DOMINATE your interview with AlgoAdvance AI