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;
    }