【堆】Leetcode 373. 查找和最小的 K 对数字【中等】

查找和最小的 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]

解题思路

• 使用最小堆（优先队列）：利用最小堆来存储当前可能的最小数对，并每次从堆中取出最小的数对。
• 初始化堆：将数组 nums1 中的前 k 个元素与 nums2 的第一个元素配对并加入堆中。由于数组是有序的，nums1[i] + nums2[0] 一定是 nums1[i] 相关的最小和。
• 扩展堆：每次从堆中取出最小的数对 (u, v)，然后将 (u, v_next)（其中 v_next 是 v 在 nums2 中的下一个元素）加入堆中，一定是 nums2[j] 相关的最小和，nums1[i]相关的最小和和nums2[j]相关的最小和都在一个最小堆中，一定能获取到当前的最小和。
• 停止条件：重复上述过程直到找到 k 个数对。

Java实现

public class KSmallestPairs {
    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;
        }

        // 最小堆，堆元素是三元组 (sum, i, j)
        PriorityQueue<int[]> minHeap = new PriorityQueue<>(Comparator.comparingInt(a -> a[0]));

        // 初始化堆，nums1 中的前 k 个元素与 nums2 的第一个元素配对
        for (int i = 0; i < Math.min(nums1.length, k); i++) {
            minHeap.offer(new int[]{nums1[i] + nums2[0], i, 0});
        }

        // 找到 k 个数对
        while (k-- > 0 && !minHeap.isEmpty()) {
            int[] current = minHeap.poll();  // 从堆中取出最小元素
            int i = current[1];  // 获取第一个数组 nums1 的索引
            int j = current[2];  // 获取第二个数组 nums2 的索引
            List<Integer> pair = new ArrayList<>();
            pair.add(nums1[i]);
            pair.add(nums2[j]);
            result.add(pair);  // 将当前数对加入结果列表

            // 如果 nums2 中有下一个元素，则将 (nums1[i], nums2[j+1]) 加入堆
            if (j + 1 < nums2.length) {
                minHeap.offer(new int[]{nums1[i] + nums2[j + 1], i, j + 1});
            }
        }

        return result;
    }

    // 测试用例
    public static void main(String[] args) {
        KSmallestPairs solution = new KSmallestPairs();
        int[] nums1 = {1, 7, 11};
        int[] nums2 = {2, 4, 6};
        int k = 3;

        List<List<Integer>> result = solution.kSmallestPairs(nums1, nums2, k);
        for (List<Integer> pair : result) {
            System.out.println(pair);
        }
        // 期望输出: [1, 2], [1, 4], [1, 6]
    }
}

时间空间复杂度

• 时间复杂度：初始化堆的时间复杂度是 O(klogk)，因为将最多 k 个元素加入堆。
  在最坏情况下，堆中最多包含 k 个元素，每次操作（插入和删除）需要 O(logk) 时间。因此，总体时间复杂度为 O(klogk+klogk)=O(klogk)。

• 空间复杂度：最小堆最多包含 k 个元素，因此空间复杂度是 O(k)。