/ Array

Find K Pairs with Smallest Sum Java Solution

Problem Statement

Given two integer arrays, arr1 and arr2, find K pairs of elements such that (x,y) where x is the value from the first array and y is the value from the second array who sum is the smallest.
For example, give the input,
arr1 = {1, 3, 4}, arr2 = {2, 6, 9}, K=3
Output would be [[1,2],[3,2],[4,2]]


Solution Explanation

Given that there are two integer arrays each with size 'm' and 'n', the number of different pair of elements would be m*n.

From m*n number of pair of elements, we need to find K pairs of elements whose sum is the lowest.

The keyword here is 'K' min number of pairs, which says we have to use a heap and since it's min number of pairs, we use a min heap. We will use a priority queue, with an integer array as the input element and a comparator which orders the elements based on the difference in the sum of the array elements.

//Order elements based on the sum of the first 2 array elements.
PriorityQueue<int[]> queue = new PriorityQueue<>((a,b)-> a[0]+a[1] -b[0] -b[1]);

Since the array elements are already in sorted order, we will add the combination of all the elements of first array and the first element of second array initially. This will allow us to limit the number of elements in the priority queue to K.

 //We add the third element as an index to which element from the second array we are pointing to   
    for(int i=0;i<arr1.length && i<k;i++) queue.offer(new int[]{arr1[i], arr2[0], 0});

Now we can process the elements in the queue until the queue is empty or K is greater than 0 and add each pair to the result set. After we add each element to the result set, we add a new pair by using the third element as index in the integer array pair.

Below is the full java implementation of the find K pairs with smallest sum:

class FindKSmallestPairs {
    public static void main(String[] args){
        int[] nums1 = {1, 3, 4};
        int[] nums2 = {2, 6, 9};
        System.out.println(new FindKSmallestPairs().kSmallestPairs(nums1, nums2, 3));
    }
    public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        int nlen1 = nums1.length, nlen2 = nums2.length;
        List<int[]> res = new ArrayList<>();
        if(nlen1==0 || nlen2==0 || k==0) return res;
        PriorityQueue<int[]> queue = new PriorityQueue<>((a, b) -> a[0]+a[1]-b[0]-b[1]);
        for(int i=0; i<nlen1 && i<k ; i++) queue.offer(new int[] {nums1[i], nums2[0], 0});
        while(k-- > 0 && !queue.isEmpty()){
            int[] cur = queue.poll();
            res.add(new int[]{cur[0],cur[1]});
            if(cur[2]==nums2.length-1) continue;
            queue.offer(new int[]{cur[0], nums2[cur[2]+1], cur[2]+1});
        }
        return res;
    }
}
Time Complexity

The time complexity is O(K logK).