373.查找和最小的 K 对数字

题目:373.查找和最小的 K 对数字

给定两个以 非递减顺序排列 的整数数组 nums1 和 nums2 , 以及一个整数 k 。

定义一对值 (u,v)，其中第一个元素来自 nums1，第二个元素来自 nums2 。

请找到和最小的 k 个数对 (u1,v1), (u2,v2) … (uk,vk) 。

• 示例 1：

输入: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
输出: [1,2],[1,4],[1,6]
解释: 返回序列中的前 3 对数：
     [1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]

• 示例 2：

输入: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
输出: [1,1],[1,1]
解释: 返回序列中的前 2 对数：
     [1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]

• 提示：

1 <= nums1.length, nums2.length <= 10^5
-10^9 <= nums1[i], nums2[i] <= 10^9
nums1 和 nums2 均为 升序排列
1 <= k <= 10^4
k <= nums1.length * nums2.length

思路

这个问题要求找到两个数组中和最小的 k 个数对 (u, v)。我们可以通过最小堆来解决这个问题，利用堆的特性来高效地找到最小的 k 个和的数对。

1. 首先，nums1 和 nums2 都是按非递减顺序排列的。利用这个条件，结果中最小的组合肯定是 nums1[0] + nums2[0]。但是，次小的是什么呢？ 是 nums1[0] + nums2[1] 还是 nums1[1] + nums2[0] 呢？我们不知道。所以，我们需要设计一种比较方式使我们能知道上述比较的结果，使用最小堆就可以解决。

2. 我们可以使用最小堆来存储这些数对的索引，并按照数组对和由小到大排序。

3. 将 nums1 中的前 k 个元素的索引与 nums2 第一个元素的索引配对，并将它们加入最小堆中。即入队的元素有 [0, 0]、[1, 0]、[2, 0]、[3, 0]、......。

4. 每次从堆中取出一个最小的数对，将其加入结果集中。然后，将下一个潜在的最小数对（即 nums1 中相同元素与 nums2 的下一个元素的组合）加入堆中。

   • 首次弹出的肯定是 [0, 0]，再把 [0, 1] 入队；
   • 这样就可以通过优先级队列比较 [0, 1] 和 [1, 0] 的结果，再弹出较小者；

5. 重复上述过程，直到找到 k 个最小和的数对。

• 时间复杂度:O(k)
• 空间复杂度:O(k)

代码

public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
    List<List<Integer>> result = new LinkedList<>();
    // 最小堆，存储索引，并按照数组对和由小到大排序
    PriorityQueue<int[]> minHeap = new PriorityQueue<>(Comparator.comparingInt(a -> (nums1[a[0]] + nums2[a[1]])));
    // 初始化堆：将 nums1 中的前 k 个元素的索引与 nums2 第一个元素的索引配对
    for (int i = 0; i < k && i < nums1.length; i++) {
        minHeap.offer(new int[]{i, 0});
    }
    while (!minHeap.isEmpty() && k > 0) {
        int[] current = minHeap.poll();
        int i = current[0];
        int j = current[1];
        result.add(Arrays.asList(nums1[i], nums2[j]));
        // 将当前数对中，nums1[i] 和 nums2[j+1]的组合加入堆
        if (j + 1 < nums2.length) {
            minHeap.offer(new int[]{i, j + 1});
        }
        k--;
    }
    return result;
}