373. Find K Pairs with Smallest Sums – Leetcode Solutions

Example 1:

Input: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
Output: [[1,2],[1,4],[1,6]]
Explanation: The first 3 pairs are returned from the sequence: [1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]

Example 2:

Input: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
Output: [[1,1],[1,1]]
Explanation: The first 2 pairs are returned from the sequence: [1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        # Edge case: if any of the lists is empty, return an empty list as no pairs can be formed
        if not nums1 or not nums2:
            return []
        
        # Min-heap to store the pairs with their sums
        min_heap = []
        result = []

        # Initialize the heap with pairs (nums1[i], nums2[0]) for i in range(len(nums1))
        # Only go up to min(k, len(nums1)) because we only need at most k pairs
        for i in range(min(k, len(nums1))):
            heapq.heappush(min_heap, (nums1[i] + nums2[0], i, 0))
        
        # Extract the minimum pairs from the heap up to k times
        while min_heap and len(result) < k:
            # Get the smallest sum pair from the heap
            curr_sum, i, j = heapq.heappop(min_heap)
            result.append([nums1[i], nums2[j]])
            
            # If there's a next element in nums2, add the pair (nums1[i], nums2[j+1]) to the heap
            if j + 1 < len(nums2):
                heapq.heappush(min_heap, (nums1[i] + nums2[j+1], i, j + 1))

        return result

class Solution {
public:
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        // Result vector to store the k smallest pairs
        vector<vector<int>> result;
        
        // Edge case: if either array is empty, return an empty result
        if (nums1.empty() || nums2.empty() || k == 0) {
            return result;
        }
        
        // Min-heap to store pairs with their sum, 
        // the heap elements are tuples: (sum, index in nums1, index in nums2)
        auto compare = [](const tuple<int, int, int>& a, const tuple<int, int, int>& b) {
            return get<0>(a) > get<0>(b); // min-heap based on pair sum
        };
        priority_queue<tuple<int, int, int>, vector<tuple<int, int, int>>, decltype(compare)> minHeap(compare);
        
        // Initialize the heap with pairs from the first element of nums1 and each element in nums2
        for (int j = 0; j < nums2.size() && j < k; ++j) {
            minHeap.emplace(nums1[0] + nums2[j], 0, j);
        }
        
        // Extract the smallest pairs until we have k pairs or the heap is empty
        while (k-- > 0 && !minHeap.empty()) {
            // Get the smallest sum pair from the heap
            auto [sum, i, j] = minHeap.top();
            minHeap.pop();
            
            // Add this pair to the result
            result.push_back({nums1[i], nums2[j]});
            
            // If there's a next element in nums1, add the new pair to the heap
            if (i + 1 < nums1.size()) {
                minHeap.emplace(nums1[i + 1] + nums2[j], i + 1, j);
            }
        }
        
        return result;
    }
};

var kSmallestPairs = function(nums1, nums2, k) {
    // Edge case: if any of the arrays is empty, return an empty array
    if (nums1.length === 0 || nums2.length === 0) return [];
    
    // Initialize a min-heap to store pairs in the form of [sum, i, j]
    // where `sum` is the sum of nums1[i] + nums2[j], i is the index in nums1, and j is the index in nums2
    let minHeap = new MinPriorityQueue({ priority: (x) => x[0] });
    
    // Push initial pairs (nums1[i], nums2[0]) into the min-heap for i < nums1.length
    for (let i = 0; i < Math.min(nums1.length, k); i++) {
        minHeap.enqueue([nums1[i] + nums2[0], i, 0]);  // Sum of pair, index in nums1, index in nums2
    }

    // Initialize result array
    let result = [];
    
    // Extract `k` smallest pairs from the heap
    while (k > 0 && !minHeap.isEmpty()) {
        let [sum, i, j] = minHeap.dequeue().element; // Get the smallest sum pair from heap
        result.push([nums1[i], nums2[j]]);  // Add the corresponding pair to result
        
        // Move to the next element in nums2 for the current element in nums1
        if (j + 1 < nums2.length) {
            minHeap.enqueue([nums1[i] + nums2[j + 1], i, j + 1]);
        }
        
        // Decrement k as we've added one pair to the result
        k--;
    }

    return result;
};

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        // Result list to store the k smallest pairs
        List<List<Integer>> result = new ArrayList<>();
        
        // Edge case: if either nums1 or nums2 is empty, return an empty result list
        if (nums1.length == 0 || nums2.length == 0 || k == 0) {
            return result;
        }
        
        // Min-heap to store pairs along with their sum
        PriorityQueue<int[]> minHeap = new PriorityQueue<>(
            (a, b) -> Integer.compare(a[0] + a[1], b[0] + b[1])
        );
        
        // Initialize the heap with pairs consisting of the first element in nums1 and each element in nums2
        // We only pair nums1[0] with nums2 elements initially to reduce complexity
        for (int j = 0; j < nums2.length && j < k; j++) {
            minHeap.offer(new int[] { nums1[0], nums2[j], 0 });
        }
        
        // Extract the smallest pairs up to k times
        while (k > 0 && !minHeap.isEmpty()) {
            // Get the pair with the smallest sum from the heap
            int[] current = minHeap.poll();
            int num1 = current[0];
            int num2 = current[1];
            int i = current[2];
            
            // Add this pair to the result
            result.add(Arrays.asList(num1, num2));
            
            // Decrement k since we added a new pair to the result
            k--;
            
            // If there are more elements in nums1 to pair with the current nums2 element, add the next pair to the heap
            if (i + 1 < nums1.length) {
                minHeap.offer(new int[] { nums1[i + 1], num2, i + 1 });
            }
        }
        
        return result;
    }
}