Find K Pairs with Smallest Sums

Problem

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.

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

Example 1:

Given nums1 = [1,7,11], nums2 = [2,4,6],  k = 3

Return: [1,2],[1,4],[1,6]

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:

Given nums1 = [1,1,2], nums2 = [1,2,3],  k = 2

Return: [1,1],[1,1]

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]

Example 3:

Given nums1 = [1,2], nums2 = [3],  k = 3 

Return: [1,3],[2,3]

All possible pairs are returned from the sequence:
[1,3],[2,3]

Solution

Brute Force Solution using SubList(): O(n * m) time

public class Solution{
    public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<int[]> res = new ArrayList<>();
        for (int i = 0; i < nums1.length; i++) {
            for (int j = 0; j < nums2.length; j++) {
                res.add(new int[] {nums1[i], nums2[j]});
            }
        }
        Collections.sort(res, (a, b) -> a[0] + a[1] - b[0] - b[1]);
        if (res.size() < k) return res;
        return res.subList(0, k); //Must indicate both starting and ending index for subList() method

    }
}

PriorityQueue Solution with Index: O(klogk) time

public class Solution {
    public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<int[]> res = new ArrayList<>();
        if (nums1.length == 0 || nums2.length == 0 || k == 0) return res;
        PriorityQueue<int[]> q = new PriorityQueue<>((a,b) -> a[0] + a[1] - b[0] - b[1]);
        for (int i = 0; i < nums1.length && i < k; i++) q.offer(new int[] {nums1[i], nums2[0], 0}); //The index is always 0 here
        while (k-- > 0 && !q.isEmpty()) { //Do not forget to check the value of k
            int[] cur = q.poll();
            res.add(new int[] {cur[0], cur[1]});
            int indexOfNums2 = cur[2];
            if (indexOfNums2 < nums2.length - 1) q.offer(new int[] {cur[0], nums2[++indexOfNums2], indexOfNums2});
        }
        return res;
    } 
}

Analysis

The first solution collects all possible pairs with two nested loops
Then we sort the collection and use subList(start, ent) to return correct result

The second solution makes use of PriorityQueue provided in Java library
We first construct our queue passing the lambada comparator new PriorityQueue<>((a,b) -> a[0] + a[1] - b[0] - b[1])
Then we loop through the first given array nums1 (of course you can loop the nums2, it does not really matter)
We construct an int array with every element from looping array nums1[i], nums2[0], and index 0
The idea behind this is that we know for every number in nums1, the smallest sum pair is with the first element in nums2, because they are both sorted in increasing order
Then we do a while loop as long as k-- > 0 and the queue is not empty
We pop the array from queue, and then construct the pair and add the array to our result
After that, if index of nums2, which is the third element of array from our queue, is not the last one, we construct a new pair with next element from nums2 and add it to our queue
Since we define our PriorityQueue with the sum, we guarantee we get the k smallest pairs ;)