Leetcode 373 查找和最小的k对数字

题目

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

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

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

解题思路

  优先队列。对于一对数 (a1, b1) ，离它最近且比它大的数可能是 (a1 + 1, b1) 或者 (a1, b1 + 1) 。那么对于已经选择的数可以推出可能的数，我们可以将可能的数维护一个升序的优先队列，每次取队列顶端的元素。由于枚举的时候会出现重复的情况，所以我们可以将 nums1 的前 k 个数对 (0, 0),(1, 0),(k - 1, 0) 先添加到优先队列中，每次将每对数的 (a1, b1 + 1) 入队即可。

代码

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        Queue<int[]> pq = new PriorityQueue<>(k, new Comparator<int[]>() {
            @Override
            public int compare(int[] o1, int[] o2) {
                return nums1[o1[0]] + nums2[o1[1]] - nums1[o2[0]] - nums2[o2[1]];
            }
        });

        int lengthNum1 = nums1.length, lengthNum2 = nums2.length;
        for (int i = 0; i < Math.min(k, lengthNum1); i++) pq.offer(new int[]{i, 0});

        List<List<Integer>> ans = new ArrayList<>();
        while (!pq.isEmpty() && k-- > 0) {
            int[] now = pq.poll();
            List<Integer> tmp = new ArrayList<>();
            tmp.add(nums1[now[0]]);
            tmp.add(nums2[now[1]]);
            ans.add(tmp);
            if (now[1] + 1 < lengthNum2) pq.offer(new int[]{now[0], now[1] + 1});
        }
        return ans;
    }
}