Find K Pairs with Smallest Sums

Problem Highlights

• 🔗 Leetcode Link: https://leetcode.com/problems/find-k-pairs-with-smallest-sums/
• 💡 Problem Difficulty: Medium
• ⏰ Time to complete: 10 to 12 mins
• 🛠️ Topics: Heap
• 🗒️ Similar Questions: Kth Smallest Element in a Sorted Matrix

1: U-nderstand

Understand what the interviewer is asking for by using test cases and questions about the problem.

• Established a set (2-3) of test cases to verify their own solution later.
• Established a set (1-2) of edge cases to verify their solution handles complexities.
• Have fully understood the problem and have no clarifying questions.
• Have you verified any Time/Space Constraints for this problem?

Be sure that you clarify the input and output parameters of the problem:

• What kind of heap can we use?
  • For finding the k smallest pairs, we will need a max-heap of size k. The idea is to take all combinations of elements from both lists, and keep adding them to a max-heap until it becomes full.
• What do we compare once the max heap is full?
  • Once the max heap is full, we compare the sum of the new pair with the top element in max heap.
• When do we pop and add our new element?
  • Only if the new sum is lesser than the sum of the max pair in the heap

Run through a set of example cases:

HAPPY CASE
Input: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
Output: [[1,2],[1,4],[1,6]]

Input: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
Output: [[1,1],[1,1]]

EDGE CASE
Input: nums1 = [1,2], nums2 = [3], k = 3
Output: [[1,3],[2,3]]

2: M-atch

Match what this problem looks like to known categories of problems, e.g. Linked List or Dynamic Programming, and strategies or patterns in those categories.

• Heaps: In the heap data structure, we assign key-value or weight to every node of the tree. Now, the root node key value is compared with the children’s nodes and then the tree is arranged accordingly into two categories i.e., max-heap and min-heap. Heap data structure is basically used as a heapsort algorithm to sort the elements in an array or a list.
• Brute force: We could build every pair from the two lists and sort all of them by sum. Would that help solve the problem? It produces the right answer, but it materializes all n1 * n2 pairs; a heap of size k lets us keep only the k best candidates seen so far.
• Sorted input: Both lists are given in ascending order. Would that help us solve this problem? Yes — once a pair's sum is no smaller than the largest sum in a full heap, every later pair in that row is at least as large, so we can stop early instead of examining every combination.

3: P-lan

Plan the solution with appropriate visualizations and pseudocode.

General Idea: For finding the k smallest pairs, use a max-heap of size k. The idea is to take all combinations of elements from both lists, and keep adding them to a max-heap until it becomes full.

1) Initialize a heap to store the sum of pairs and their respective indexes.
2) Initially, we start making pairs by pairing only the first element of the second list with each element of the first list. We push the pairs onto a max-heap, sorted by the sum of each pair.
3) Use another loop to pop the smallest pair from the max-heap, noting the sum of the pair and the list indexes of each element, and add the pair to a result list.
4) To make new pairs, we move forward in the second list and pair the next element in it with each element of the first list, pushing each pair on the max-heap.
5) We keep pushing and popping pairs from the max-heap until we have collected the required k smallest pairs in the result list.

⚠️ Common Mistakes

• Is it possible for nums1[0] + nums2[k+1] can be a candidate? This cannot be, because nums1[0] + nums2[0....k] is always smaller than nums1[0] + nums2[k+1]. Each time after we pick the pair with min sum, we put the new pair with the second index +1. ie, pick (0,0), we put back (0,1). Therefore, the heap always maintains at most min(k, len(nums1)) elements.

4: I-mplement

Implement the code to solve the algorithm.

import heapq
from typing import List

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        heap = []
        nums1.sort()
        nums2.sort()

        for n1 in nums1:
            # once the heap is full and this row's smallest sum can't beat the
            # largest sum in the heap, no later row can either
            if len(heap) >= k and n1 + nums2[0] > -heap[0][0]:
                break
            for n2 in nums2:
                sm = n1 + n2
                if len(heap) < k:
                    heapq.heappush(heap, (-sm, (n1, n2)))
                elif sm < -heap[0][0]:
                    heapq.heappop(heap)
                    heapq.heappush(heap, (-sm, (n1, n2)))
                else:
                    # sums only grow as n2 increases, so the rest of this row
                    # can't beat the heap's largest sum
                    break

        ls = []
        for ele in heap:
            ls.append(list(ele[1]))

        return ls

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        // we set a pq to store the minimum sum
        PriorityQueue<List<Integer>> pq = new PriorityQueue<>((i1, i2) -> ((i1.get(0) + i1.get(1)) - (i2.get(0) + i2.get(1))));

        // try first k elements in nums1 here
        for(int i = 0; i < Math.min(nums1.length, k); i++){
            pq.add(new ArrayList<Integer>(Arrays.asList(nums1[i], nums2[0], 0)));
        }
        
        List<List<Integer>> ret = new ArrayList<>();

        while(!pq.isEmpty() && k > 0){
            List<Integer> cur = pq.poll();
            ret.add(new ArrayList<>(Arrays.asList(cur.get(0), cur.get(1))));
            k--;
            
            // because we only do the traverse in nums1
            // so we sum the current smallest element in nums1
            // with the current searched smallest element in nums2
            if(cur.get(2) + 1 < nums2.length){
                pq.offer(new ArrayList<Integer>(Arrays.asList(cur.get(0), nums2[cur.get(2) + 1], cur.get(2) + 1)));
            }
        }
        
        return ret;
    }
}   

5: R-eview

Review the code by running specific example(s) and recording values (watchlist) of your code's variables along the way.

• Trace through your code with an input to check for the expected output
• Catch possible edge cases and off-by-one errors

6: E-valuate

Evaluate the performance of your algorithm and state any strong/weak or future potential work.

Let n1 and n2 be the lengths of nums1 and nums2.

• Time Complexity: Python: O(n1 * n2 * log k) in the worst case — the nested loops can examine every pair, each with an O(log k) heap operation, though the early breaks prune much of this in practice. Java: O(k * log(min(k, n1))), since its heap never holds more than min(k, n1) entries and we pop/push at most k times.
• Space Complexity: O(k) for the Python max-heap and O(min(k, n1)) for the Java min-heap, plus the O(k) output list