Leetcode: Find K Pairs with Smallest Sums

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]

 1 public class Solution {
 2     public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
 3         PriorityQueue<int[]> que = new PriorityQueue<>((a,b)->a[0]+a[1]-b[0]-b[1]);
 4         List<int[]> res = new ArrayList<>();
 5         if(nums1.length==0 || nums2.length==0 || k==0) return res;
 6         for(int i=0; i<nums1.length && i<k; i++) que.offer(new int[]{nums1[i], nums2[0], 0});
 7         while(k-- > 0 && !que.isEmpty()){
 8             int[] cur = que.poll();
 9             res.add(new int[]{cur[0], cur[1]});
10             if(cur[2] == nums2.length-1) continue;
11             que.offer(new int[]{cur[0],nums2[cur[2]+1], cur[2]+1});
12         }
13         return res;
14     }
15 }

 1 public class Solution {
 2     public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
 3         List<int[]> res = new ArrayList<int[]>();
 4         Comparator<int[]> comp = new Comparator<int[]>() {
 5             public int compare(int[] arr1, int[] arr2) {
 6                 int sum1 = arr1[0] + arr1[1];
 7                 int sum2 = arr2[0] + arr2[1];
 8                 if (sum1 == sum2) return arr1[0]-arr2[0];
 9                 else return sum1 - sum2;
10             }
11         };
12         
13         PriorityQueue<int[]> minSumHeap = new PriorityQueue<int[]>(11, comp);
14         for (int i=0; i<nums1.length; i++) {
15             for (int j=0; j<nums2.length; j++) {
16                 int[] arr = new int[2];
17                 arr[0] = nums1[i];
18                 arr[1] = nums2[j];
19                 minSumHeap.offer(arr);
20             }
21         }
22         int count = 0;
23         while (count<k && !minSumHeap.isEmpty()) {
24             res.add(minSumHeap.poll());
25             count++;
26         }
27         return res;
28     }
29 }