【LeetCode】373. Find K Pairs with Smallest Sums

373. Find K Pairs with Smallest Sums

You are given two integer arrays nums1 and nums2 sorted in ascending order and an integer k.

Define a pair (u,v) which consists of one element from the first array and one element from the second array.

Find the k pairs (u1,v1),(u2,v2) …(uk,vk) with the smallest sums.

Example 1:

Given nums1 = [1,7,11], nums2 = [2,4,6],  k = 3

Return: [1,2],[1,4],[1,6]

The first 3 pairs are returned from the sequence:
[1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]

Example 2:

Given nums1 = [1,1,2], nums2 = [1,2,3],  k = 2

Return: [1,1],[1,1]

The first 2 pairs are returned from the sequence:
[1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]

Example 3:

Given nums1 = [1,2], nums2 = [3],  k = 3 

Return: [1,3],[2,3]

All possible pairs are returned from the sequence:
[1,3],[2,3]

解析：

从两个排序数组中分别各取出一个数组成一对序列，题目的要求是取出组成序列和最小的前K个序列（u1+v1）最小的。

这道题可以用两种解法：

• 暴力法（hash+list）

  采用hashMap的key存储（u1+v1），value使用一个List来存放序列。这里必须使用List，因为可能存在相同key，这样会将原始的序列给覆盖，所以遇见相同的key我们使用List来链接到后面。然后遍历map,针对每个key的list进行遍历，直至取出的元素大小为k结束。

• 最小堆

  利用最小堆的特性，顶端的元素最小，我们可以构建一个最小堆，最小堆我们存放两个数组的index索引，在构建的过程中，我们就可以吧索引存好，这样我们直接遍历前k个元素，通过所用来获取原数组的值即可。

  我们还可以边构建，边取值，题目中的数组是从小到大排序的，这里给了个提示 num1[0]+num2[0]的值一定是最小的，所以我们先将这个值扔进堆中，然后开始遍历下一个元素，下一个值一定是从num1[1]+num2[0]和num1[0]+num2[1]之中取出最小的那个，依次循环，直至取出的元素为K。

代码：

暴力：

class Solution {
    public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        Map<Integer, List<int[]>> map = new TreeMap<>(Integer::compareTo);
        for (int num1 : nums1) {
            for (int num2 : nums2) {
                if (map.containsKey(num1 + num2)) {
                    map.get(num1 + num2).add(new int[]{num1, num2});
                } else {
                    List list = new ArrayList<>();
                    list.add(new int[]{num1, num2});
                    map.put(num1 + num2, list);
                }

            }
        }
        List<int[]> ret = new ArrayList<>();


        for (Map.Entry<Integer, List<int[]>> entry : map.entrySet()) {
            List<int[]> list = entry.getValue();
            if (k >= list.size()) {
                ret.addAll(list);
                k -= list.size();
            } else if (k > 0) {
                ret.addAll(list.stream().limit(k).collect(Collectors.toList()));
                k = 0;
            }
        }
        return ret;
    }
}

最小堆：

public static List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<int[]> ret = new ArrayList<>();
        if(nums1.length == 0 || nums2.length == 0){
            return ret;
        }
        boolean visit[][] = new boolean[nums1.length][nums2.length];
        PriorityQueue<int[]> queue = new PriorityQueue<>(new Comparator<int[]>() {
            @Override
            public int compare(int[] o1, int[] o2) {
                return nums1[o1[0]] + nums2[o1[1]] - nums1[o2[0]] - nums2[o2[1]];
            }
        });

        queue.add(new int[]{0, 0});
        visit[0][0] = true;

        while (!queue.isEmpty() && ret.size() < Math.min(k, nums1.length * nums2.length)) {

            final int[] index = queue.poll();
            ret.add(new int[]{nums1[index[0]], nums2[index[1]]});

            if (index[0] + 1 < nums1.length && !visit[index[0] + 1][index[1]]) {
                visit[index[0] + 1][index[1]] = true;
                queue.add(new int[]{index[0] + 1, index[1]});
            }

            if (index[1] + 1 < nums2.length && !visit[index[0]][index[1] + 1]) {
                visit[index[0]][index[1] + 1] = true;
                queue.add(new int[]{index[0], index[1] + 1});
            }
        }

        return ret;
    }