Solution.java

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.PriorityQueue;
import java.util.TreeSet;

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<List<Integer>> result = new ArrayList<>();
        if (nums1.length == 0 || nums2.length == 0 || k == 0) {
            return result;
        }
        PriorityQueue<int[]> Q = new PriorityQueue<>((a, b) -> nums1[a[0]] + nums2[a[1]] - nums1[b[0]] - nums2[b[1]]);
        TreeSet<int[]> added = new TreeSet<>((a, b) -> {
            if (a[0] != b[0]) {
                return a[0] - b[0];
            }
            return a[1] - b[1];
        });

        Q.add(new int[]{0, 0});
        added.add(new int[]{0, 0});
        while (result.size() < k && !Q.isEmpty()) {
            int[] ans = Q.poll();
            int i = ans[0];
            int j = ans[1];
            result.add(Arrays.asList(nums1[i], nums2[j]));
            if (i + 1 < nums1.length && !added.contains(new int[]{i + 1, j})) {
                Q.add(new int[]{i + 1, j});
                added.add(new int[]{i + 1, j});
            }
            if (j + 1 < nums2.length && !added.contains(new int[]{i, j + 1})) {
                Q.add(new int[]{i, j + 1});
                added.add(new int[]{i, j + 1});
            }
        }
        return result;
    }
}