优先队列 -leetcode-373. 查找和最小的 K 对数字

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

题目描述

给定两个以 升序排列 的整数数组 nums1 和 nums2 , 以及一个整数 k 。
定义一对值 (u,v)，其中第一个元素来自 nums1，第二个元素来自 nums2 。
请找到和最小的 k 个数对 (u1,v1), (u2,v2) … (uk,vk) 。

优先队列

题目数据范围过大 $10^5$，暴力算法算出两两结果，会 tle

思路：彤哥题解

1. 使用堆保存 [index1, index2]，即当前最小数对为：{nums1[index1], nums2[index2]}
2. 将 [0, 0], [1, 0], …, [n1, 0]添加到 heap中；
3. 从当前 heap中弹出堆顶元素为 arr，并保存 (nums1[arr[0]], nums2[arr[1]]))
4. 将 index2后移一位；
5. 最多从 heap弹出 k次.

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        // 保存 [index1, index2]
        PriorityQueue<int[]> heap = new PriorityQueue<>((a, b) -> (nums1[a[0]] + nums2[a[1]]) - (nums1[b[0]] + nums2[b[1]]));
        // 把nums1中所有元素入堆。（优化：只需要k个）
        for (int i = 0; i < Math.min(nums1.length, k); i++) {
            heap.offer(new int[] {i, 0});
        }
        // System.out.println(heap);
        List<List<Integer>> res = new ArrayList<>();
        // 最多弹出 k次
        while (k-- > 0 && !heap.isEmpty()) { // 这里一定要判空
            int[] arr = heap.poll();
            System.out.println(Arrays.toString(arr));
            res.add(Arrays.asList(nums1[arr[0]], nums2[arr[1]]));
            // index2加一后继续入堆
            if (++arr[1] < nums2.length) {
                heap.offer(arr);
            }
        }
        return res;
    }
}