LeetCode 373. Find K Pairs with Smallest Sums

LC address: 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 (u₁,v₁),(u₂,v₂) …(u_k,v_k) 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]

Analysis:

首先对题目进行分析，为了方便，我用Ai来表示nums1[i]，用Bi来表示nums2[i]。我们可以知道[Ai, Bj]永远比[Ai+m, Bj+n]小（m+n>0）。换句话说，如果[Ai, Bj]还没有被排进list，那么[Ai+m, Bj+n]永远不可能被考虑加入list。再换句话说，在任意一个时刻，对于每一个Ai（Bi），最多只需要考虑一种组合。根据这个分析的结果，我的想法是，控制一个集合S，里面包含所有当前会被考虑排进list的组合，每次在其中选出和最小的，并更新这个集合。

为了不那么抽象，我决定用图片辅助说明，不妨有A={A1, A2, A3, A4}以及B={B1, B2, B3, B4}。假设在某一时刻，list={[A1, B1], [A2, B1], [A3, B1], [A1, B2]}。下图中黑色线条代表已经加入list的组合，红色代表当前需要被考虑的组合，也就是S={[A1,B3], [A2, B2], [A1, B3]}。

那么这个组合怎么得到呢？其实规律很简单，对于每一个Ai，如果在list中和Ai组合的最大数字为Bj，那么只需要考虑[Ai, Bj+1]；如果[Ai, B1]进了list，而[Ai + 1, B1]还没有进list，那么需要考虑[Ai + 1, B1]。

根绝这个思路，那么可以设计算法，既然每次我们都需要找S中最小的组合来加入list，很容易就想到用一个PriorityQueue来储存和更新S。到了这一步，基本也就解决了这道题，只需要再考虑edge cases就可以了。

不过依我对于LC test cases的尿性的了解，以及我多年来谨慎小心的作风，我觉得LC很可能在一些小的edge cases来阴我。对于这道题，容易想到，如果我们对于两个int求和（或者做差），可能会有overflow（underflow）的问题，所以我当即为我的机智点了赞，然后哼哧哼哧地把很多地方的int换成了long。然后我就一遍AC，心情大好，心想毕竟是心思缜密的我啊！但是我是一个强迫症，想看看不改成int会有什么结果，然而依旧AC！(╯‵□′)╯︵┻━┻  LC你太让我失望了，test cases写得这么马虎，怎么对得起我付的钱？！还钱！

Solution:

solution 1 : 简约版

public class Solution {
    public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        
        List<int[]> list = new ArrayList<>();
        if (nums1 == null || nums1.length == 0
            ||nums2 == null || nums2.length == 0) {
            return list;
        }
        
        PriorityQueue<int[]> pq = new PriorityQueue<>(new Comparator<int[]>() {
            public int compare(int[] pair1, int[] pair2) {
                return nums1[pair1[0]] + nums2[pair1[1]]
                    - nums1[pair2[0]] - nums2[pair2[1]];
            }
        });
        
        int[] temp = new int[2];
        pq.add(temp);
        int count = 0;
        
        while (k > 0 && !pq.isEmpty()) {
            temp = pq.poll();
            int[] res = new int[2];
            int i1 = temp[0];
            int i2 = temp[1];
            res[0] = nums1[i1];
            res[1] = nums2[i2];
            list.add(res);
            if (i1 + 1 < nums1.length) {
                temp[0] = i1 + 1;
                pq.add(temp);
            }
            if (i1 == 0 && i2 == count && count + 1 < nums2.length) {
                count += 1;
                int[] newArray = new int[2];
                newArray[1] = count;
                pq.add(newArray);
            }
            k -= 1;
        }
        return list;
    }
}

solution 2 : 深思熟虑谨小慎微然后并没有什么卵用版

public class Solution {
    public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        
        List<int[]> list = new ArrayList<>();
        if (nums1 == null || nums1.length == 0
            ||nums2 == null || nums2.length == 0) {
            return list;
        }
        
        PriorityQueue<long[]> pq = new PriorityQueue<>(new Comparator<long[]>() {
            public int compare(long[] pair1, long[] pair2) {
                if (pair1[2] == pair2[2]) {
                    return 0;
                }
                if (pair1[2] < pair2[2]) {
                    return -1;
                }
                return 1;
            }
        });
        
        long[] temp = new long[3];
        temp[2] = (long) nums1[0] + (long) nums2[0];
        pq.add(temp);
        int count = 0;
        
        while (k > 0 && !pq.isEmpty()) {
            temp = pq.poll();
            int[] res = new int[2];
            int i1 = (int) temp[0];
            int i2 = (int) temp[1];
            res[0] = nums1[i1];
            res[1] = nums2[i2];
            list.add(res);
            if (i1 + 1 < nums1.length) {
                temp[0] = i1 + 1;
                temp[2] = (long) nums1[i1 + 1] + nums2[i2];
                pq.add(temp);
            }
            if (i1 == 0 && i2 == count && count + 1 < nums2.length) {
                count += 1;
                long[] newArray = new long[3];
                newArray[0] = 0;
                newArray[1] = count;
                newArray[2] = (long) nums1[0] + nums2[count];
                pq.add(newArray);
            }
            k -= 1;
        }
        return list;
    }
}

Solution code can also be found here: https://github.com/all4win/LeetCode.git