JAVA算法学习-LeetCode 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]

---

解题思路

优先队列

利用优先队列（最小堆）来高效地找出两个数组中前 k 个最小的元素对。

下面是详细的解题思路：

1. 初始化优先队列（最小堆）：

   首先，创建一个优先队列 pq，其大小为 k，并且自定义了比较器。这个比较器根据两个数组中对应位置的元素之和来决定元素的优先级。这意味着队列中的元素将按照 nums1[o1[0]] + nums2[o1[1]] 的值从小到大排序。

2. 填充优先队列：

   遍历 nums1 数组的前 k 个元素（或者直到数组末尾，以先达到的那个为准），并将每个元素的索引与 nums2 数组的第一个元素的索引组成一对，加入到优先队列中。这样做的目的是初始化一些候选的最小元素对。

3. 提取最小元素对：

   接下来开始一个循环，最多执行 k 次，每次循环都会从优先队列中取出当前最小的元素对。这个元素对实际上是两个索引，可以用这两个索引来访问 nums1 和 nums2 中的元素，形成一个元素对。

4. 更新结果列表：

   从队列中取出元素对后，将对应的元素对添加到结果列表 ans 中。

5. 扩展候选元素对：

   如果刚刚取出的元素对中，nums2 的索引不是最后一个，就将新的元素对（即 nums1 的当前索引和 nums2 的下一个索引）加入到优先队列中。这样就在保持 nums1 的索引不变的情况下，探索 nums2 中下一个可能的更小元素对。

6. 重复步骤3至5：

   继续这个过程，直到达到 k 次迭代或者优先队列为空。在这个过程中，我们始终保持优先队列中包含的是当前已知的最小元素对。

7. 返回结果：

   最后，将结果列表 ans 返回，这个列表包含了两个数组中前 k 个最小的元素对。

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        // 创建一个大小为 k 的最小堆，用于存储元素对的索引
        // 比较器定义了堆中元素对的排序规则，基于两个数组的对应元素之和
        PriorityQueue<int[]> pq = new PriorityQueue<>(k, (o1, o2) -> {
            return nums1[o1[0]] + nums2[o1[1]] - nums1[o2[0]] - nums2[o2[1]];
        });
        
        // 初始化结果列表
        List<List<Integer>> ans = new ArrayList<>();
        
        // 获取两个数组的长度
        int m = nums1.length;
        int n = nums2.length;
        
        // 将 nums1 的前 k 个元素与 nums2 的第一个元素组成的索引入队
        for (int i = 0; i < Math.min(m, k); i++) {
            pq.offer(new int[]{i,0});
        }
        
        // 循环 k 次，每次从堆中取出最小的元素对
        while (k-- > 0 && !pq.isEmpty()) {
            // 取出当前最小的元素对索引
            int[] idxPair = pq.poll();
            
            // 使用取出的索引创建一个列表，并将其添加到结果列表中
            List<Integer> list = new ArrayList<>();
            list.add(nums1[idxPair[0]]);
            list.add(nums2[idxPair[1]]);
            ans.add(list);
            
            // 如果第二个数组的下一个元素还有，就将新的索引入队
            if (idxPair[1] + 1 < n) {
                pq.offer(new int[]{idxPair[0], idxPair[1] + 1});
            }
        }
        
        // 返回最终的结果列表
        return ans;
    }
}

时间复杂度：O(n log n);时间复杂度主要取决于优先队列的操作，因为在最坏情况下，需要从队列中取出 k 个元素，每次操作都需要 log k 的时间。

空间复杂度：O(n);只需要存储 k 个元素对的索引。

---