373. Find K Pairs with Smallest Sums

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<List<Integer>> resultList = new ArrayList<>();

        // 优先队列中存储的为索引对，堆顶为和最小的元素索引对
        Queue<int[]> queue = new PriorityQueue<>((o1, o2) -> {
            int sum1 = nums1[o1[0]] + nums2[o1[1]];
            int sum2 = nums1[o2[0]] + nums2[o2[1]];
            return sum1 - sum2;
        });

        // 关键之处：将 nums1 的前 k 个元素添加进去，以便于后续处理过程中无需移动 nums1 的指针，仅需移动 nums2 的指针
        for (int i = 0; i < Math.min(k, nums1.length); i++) {
            queue.offer(new int[]{i, 0});
        }

        while (k > 0 && !queue.isEmpty()) {
            int[] indexPair = queue.poll();
            int indexOfNums1 = indexPair[0], indexOfNums2 = indexPair[1];
            resultList.add(Arrays.asList(nums1[indexOfNums1], nums2[indexOfNums2]));
            k--;

            if (indexOfNums2 + 1 < nums2.length) {
                queue.offer(new int[]{indexOfNums1, indexOfNums2 + 1});
            }
        }

        return resultList;
    }
}

References

373. Find K Pairs with Smallest Sums
剑指 Offer II 061. 和最小的 k 个数对