面试算法61：和最小的k个数对

题目

给定两个递增排序的整数数组，从两个数组中各取一个数字u和v组成一个数对（u，v），请找出和最小的k个数对。例如，输入两个数组[1，5，13，21]和[2，4，9，15]，和最小的3个数对为（1，2）、（1，4）和（2，5）。

分析

这个题目要求找出和最小的k个数对。可以用最大堆来存储这k个和最小的数对。逐一将m×n个数对添加到最大堆中。当堆中的数对的数目小于k时，直接将数对添加到堆中。如果堆中已经有k个数对，那么先要比较待添加的数对之和及堆顶的数对之和（也是堆中最大的数对之和）。如果待添加的数对之和大于或等于堆顶的数对之和，那么堆中的k个数对之和都小于或等于待添加的数对之和，因此待添加的数对可以忽略。如果待添加的数对之和小于堆顶的数对之和，那么删除堆顶的数对，并将待添加的数对添加到堆中，这样可以确保堆中存储的是和最小的k个数对。每次都是将待添加的数对与堆中和最大的数对进行比较，而这也是用最大堆的原因。
接下来考虑如何优化。题目给出的条件是输入的两个数组都是递增排序的，这个特性我们还没有用到。如果从第1个数组中选出第k+1个数字和第2个数组中的某个数字组成数对p，那么该数对之和一定不是和最小的k个数对中的一个，这是因为第1个数组中的前k个数字和第2个数组中的同一个数字组成的k个数对之和都要小于数对p之和。因此，不管输入的数组nums1有多长，最多只需要考虑前k个数字。同理，不管输入的数组nums2有多长，最多也只需要考虑前k个数字。

解

public class Test {
    public static void main(String[] args) {
        int[] nums1 = {1, 5, 13, 21};
        int[] nums2 = {2, 4, 9, 15};
        List<List<Integer>> result = kSmallestPairs(nums1, nums2, 3);
        for (List<Integer> res : result) {
            System.out.println(res);
        }
    }

    public static List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        // 最大堆
        Queue<int[]> maxHeap = new PriorityQueue<>(
            (p1, p2) -> p2[0] + p2[1] - p1[0] - p1[1]
        );
        for (int i = 0; i < Math.min(k, nums1.length); i++) {
            for (int j = 0; j < Math.min(k, nums2.length); j++) {
                if (maxHeap.size() >= k) {
                    int[] root = maxHeap.peek();
                    if (root[0] + root[1] > nums1[i] + nums2[j]) {
                        maxHeap.poll();
                        maxHeap.offer(new int[] {nums1[i], nums2[j]});
                    }
                }
                else {
                    maxHeap.offer(new int[] {nums1[i], nums2[j]});
                }
            }
        }

        List<List<Integer>> result = new LinkedList<>();
        while (!maxHeap.isEmpty()) {
            int[] vals = maxHeap.poll();
            result.add(Arrays.asList(vals[0], vals[1]));
        }

        return result;
    }
}