查找和最小的 K 对数字

优质博文：IT-BLOG-CN

题目

给定两个以 非递减顺序排列 的整数数组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 <= 105
-109 <= nums1[i], nums2[i] <= 109
nums1和nums2均为 升序排列
1 <= k <= 104
k <= nums1.length * nums2.length

代码

方法一：优先队列
对于已经按升序排列的两个数组 nums₁ ,nums₂ ，长度分别为 length₁ ,length₂，我们可以知道和最小的数对一定为 (nums₁[0],nums₂[0])，和最大的数对一定为 (nums₁[length₁−1],nums₂ [length₂−1])。本题要求找到最小的 k 个数对，最直接的办法是可以将所有的数对求出来，然后利用排序或者 TopK 解法求出最小的 k 个数对即可。实际求解时可以不用求出所有的数对，只需从所有符合待选的数对中选出最小的即可，假设当前已选的前 n 小数对的索引分别为 (a₁,b₁),(a₂,b₂),(a₃,b₃),…,(a_n,b_n)，由于两个数组都是按照升序进行排序的，则可以推出第 n+1 小的数对的索引选择范围为 (a₁ +1,b₁),(a₁,b₁+1),(a₂+1,b₂),(a₂,b₂+1),(a₃+1,b₃),(a₃,b₃+1),…,(a_n +1,b_n),(a_n,b_n+1)，假设我们利用堆的特性可以求出待选范围中最小数对的索引为 (a_i ,b_i)，同时将新的待选的数对 (a_i+1,b_i),(a_i ,b_i +1) 加入到堆中，直到我们选出 k 个数对即可。

如果我们每次都将已选的数对 (a_i,b _i) 的待选索引 (a_i+1,b_i),(a_i ,b_i+1) 加入到堆中则可能出现重复的问题，一般需要设置标记位解决去重的问题。我们可以将 nums₁的前 k 个索引数对 (0,0),(1,0),…,(k−1,0) 加入到队列中，每次从队列中取出元素 (x,y) 时，我们只需要将 nums₂的索引增加即可，这样避免了重复加入元素的问题。

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int 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;
        for (int i = 0; i < Math.min(m, k); i++) {
            pq.offer(new int[]{i,0});
        }
        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(klogk)，其中 k 是选择的数对的数目。优先队列中最多只保存 k 个元素，每次压入新的元素队列进行调整的时间复杂度为 logk，入队操作一共有 2k 次, 一共需要从队列中弹出 k 个数据。

空间复杂度： O(k)。优先队列中最多只保存 k 个元素。

方法二：二分查找
利用二分查找找到第 k 小的数对和为 pairSum。利用滑动窗口即可计算出两个数组中满足数对和小于等于目标值 target 的数对有多少个，我们找到最小的 target 且满足小于等于它的数对数目刚好大于等于 k 即为目标值 pairSum，然后在数组中找到最小的 k 个数对满足数对和小于等于 pairSum。

由于题目中数组 nums ₁,nums₂中的元素存在重复，因此我们不能简单的利用滑动窗口找到所有满足小于等于 pairSum 的数对。因为存在小于等于 pairSum 的数对和的数目大于 k，因此数对和等于 pairSum 的数对不一定就属于最小的 k 个数对。
首先利用滑动窗口找到所有小于 pairSum 的数对，假设数对和小于 pairSum 的数目为 x 个，然后再利用二分查找在数组中找到 k−x 个和等于 pairSum 的数对即可。

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        int m = nums1.length;
        int n = nums2.length;

        /*二分查找第 k 小的数对和的大小*/
        int left = nums1[0] + nums2[0];
        int right = nums1[m - 1] + nums2[n - 1];
        int pairSum = right;
        while (left <= right) {
            int mid = left + ((right - left) >> 1);
            long cnt = 0;
            int start = 0;
            int end = n - 1;
            while (start < m && end >= 0) {
                if (nums1[start] + nums2[end] > mid) {
                    end--;
                } else {
                    cnt += end + 1;
                    start++;
                }
            }
            if (cnt < k) {
                left = mid + 1;
            } else {
                pairSum = mid;
                right = mid - 1;
            }
        }

        List<List<Integer>> ans = new ArrayList<>();
        int pos = n - 1;
        /*找到小于目标值 pairSum 的数对*/
        for (int i = 0; i < m; i++) {
            while (pos >= 0 && nums1[i] + nums2[pos] >= pairSum) {
                pos--;
            }
            for (int j = 0; j <= pos && k > 0; j++, k--) {
                List<Integer> list = new ArrayList<>();
                list.add(nums1[i]);
                list.add(nums2[j]);
                ans.add(list);
            }
        }

        /*找到等于目标值 pairSum 的数对*/
        pos = n - 1;
        for (int i = 0; i < m && k > 0; i++) {
            int start1 = i;
            while (i < m - 1 && nums1[i] == nums1[i + 1]) {
                i++;
            }
            while (pos >= 0 && nums1[i] + nums2[pos] > pairSum) {
                pos--;
            }
            int start2 = pos;
            while (pos > 0 && nums2[pos] == nums2[pos - 1]) {
                pos--;
            }
            if (nums1[i] + nums2[pos] != pairSum) {
                continue;
            }
            int count = (int) Math.min(k, (long) (i - start1 + 1) * (start2 - pos + 1));
            for (int j = 0; j < count && k > 0; j++, k--) {
                List<Integer> list = new ArrayList<>();
                list.add(nums1[i]);
                list.add(nums2[pos]);
                ans.add(list);
            }
        }
        return ans;
    }
}

时间复杂度： O(k+(m+n)×log(diff(nums₁)+diff(nums₂)))，其中 m,n 表示数组 nums₁,nums₂的长度，diff(arr) 表示数组 arr 中最大元素与最小元素之差，diff(nums₁)=max(nums₁)−min(nums₁),diff(nums₂)=max(nums₂)−min(nums₂))。我们利用二分查找找到满足要求的数对和的时间复杂度为 (m+n)×log(diff(nums₁)+diff(nums₂))，我们利用滑动窗口找到小于等于目标值的 k 个数对的时间复杂度为 O(2×(k+m+n))，所以总的时间复杂度 O(2×(k+m+n)+(m+n)×log(diff(nums₁)+diff(nums₂)))=k+(m+n)×log(diff(nums₁)+diff(nums₂))。

空间复杂度： O(1)，除了函数返回值以外，不需要额外的存储空间。