Find K Pairs with Smallest Sums LeetCode solution in Java

1
class Solution {
2
  public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
3
    PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) -> (a[0] + a[1]) - (b[0] + b[1]));
4
    for (int i = 0; i < nums1.length && i < k; i++) {
5
      pq.add(new int[] {nums1[i], nums2[0], 0});
6
    }
7

8
    List<List<Integer>> res = new ArrayList<>();
9
    for (int i = 0; i < k && !pq.isEmpty(); i++) {
10
      int[] curr = pq.poll();
11
      res.add(Arrays.asList(curr[0], curr[1]));
12
      int idx2 = curr[2];
13
      if (idx2 < nums2.length - 1) {
14
        pq.add(new int[] {curr[0], nums2[idx2 + 1], idx2 + 1});
15
      }
16
    }
17
    return res;
18
  }
19
}