剑指 Offer II 061. 和最小的 k 个数对（373. 查找和最小的 K 对数字）

题目：

 

 

 

思路：

【1】常规的模拟方式（虽然使用break进行了截断，但是耗费的时间还是很久）

【2】对常规方式进行优化（尽量如何减少塞入比对的元素）

【3】二分的方式进行优化

代码展示：

【1】常规的模拟方式

//时间90 ms 击败 9.43%
//内存57.1 MB 击败 26.54%
class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        int n1 = nums1.length, n2 = nums2.length;
        PriorityQueue<int[]> que = new PriorityQueue<>((a, b) -> (b[0]+b[1]) - (a[0]+a[1]));
        int[] tem;
        for (int i = 0; i < n1; i++){
            for (int j = 0; j < n2; j++){
                if (k-- > 0){
                    que.add(new int[]{nums1[i],nums2[j]});
                    continue;
                }else {
                    tem = que.peek();
                    if (tem != null && (tem[0]+tem[1]) > (nums1[i]+nums2[j])){
                        que.poll();
                        que.add(new int[]{nums1[i],nums2[j]});
                    }else {
                        // 这种情况下nums1[i]+nums2[j]要是比堆顶元素要大，那么后面的nums1[i]+nums2[j+1]也是会比堆顶元素要大
                        // 所以直接break跳出循环不用做后面的操作了
                        // 直接逃过 j++的这一层循环
                        break;
                    }
                }
            }
            // 这种情况下nums1[i]+nums2[0]要是比堆顶元素要大，那么后面的nums1[i+1]+nums2[0]也是会比堆顶元素要大
            // 所以直接break跳出循环不用做后面的操作了
            if (que.peek() != null && (que.peek()[0]+que.peek()[1]) < (nums1[i]+nums2[0])){
                //直接逃过 i++的这一层循环
                break;
            }
        }
        List<List<Integer>> res = new ArrayList<>();
        while (!que.isEmpty()){
            tem = que.poll();
            ArrayList<Integer> an = new ArrayList<Integer>();
            an.add(tem[0]);
            an.add(tem[1]);
            res.add(an);
        }
        return res;
    }
}

 

【2】对常规方式进行优化（尽量如何减少塞入比对的元素）

//时间29 ms 击败 68.28%
//内存56.6 MB 击败 38.14%
class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        //构建大顶堆队列
        PriorityQueue<int[]> pq = new PriorityQueue<>(k, (o1, o2)->(nums1[o1[0]] + nums2[o1[1]]) - (nums1[o2[0]] + nums2[o2[1]]));
        List<List<Integer>> ans = new ArrayList<>();
        int m = nums1.length , n = nums2.length;
        // 首先当num2取第一位，然后nums1取符合的长度，后面的直接就可以放弃了
        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);
            // 这里是由于第一个必然是（0,0）的下标
            // 所以下一个最小的必然是（1,0）或（0,1），由于（1,0）已经有了所以必然是加（0,1）
            // 如果下一个是（0,1），那么 下一个最小的必然是（1,1）或（0,2），
            // 但是其实只需要加（0,2）就行了，为什么呢？因为（1,1）的前面就是（1,0），
            // 所以（1,0）没有加进去的时候就必然轮不到（1,1）
            if (idxPair[1] + 1 < n) {
                pq.offer(new int[]{idxPair[0], idxPair[1] + 1});
            }
        }

        return ans;
    }
}

 

【3】二分的方式进行优化

//时间20 ms 击败 94.32%
//内存56.2 MB 击败 49.83%
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;
    }
}