JunyiCode
5/4/2020 - 5:18 PM

373. Find K Pairs with Smallest Sums

/*
    Note:
        offer initial pairs {num1, num2, index_of_num2}
        ⭐️k--;   //cannot be after: if(cur.get(2) == nums2.length - 1)  continue;

*/

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        //ascending order
        Comparator<List<Integer>> cmp = new Comparator<List<Integer>>() {
            @Override
            public int compare(List<Integer> a, List<Integer> b) {
                return a.get(0) + a.get(1) - b.get(0) - b.get(1);
            }
        };
        Queue<List<Integer>> queue = new PriorityQueue(cmp);
        // Queue<List<Integer>> queue = new PriorityQueue<>((a, b) -> a.get(0) + a.get(1) - b.get(0) - b.get(1));
        List<List<Integer>> res = new ArrayList<>();
        if(nums1 == null || nums2 == null || nums1.length * nums2.length * k == 0) 
            return res;
        for(int i = 0; i < nums1.length && i < k; i++) 
            queue.offer(new ArrayList<Integer>(Arrays.asList(nums1[i], nums2[0], 0)));
        while(k > 0 && !queue.isEmpty()) {
            k--;  //cannot be after: if(cur.get(2) == nums2.length - 1)  continue;
            List<Integer> cur = queue.poll();
            res.add(new ArrayList<Integer>(Arrays.asList(cur.get(0), cur.get(1))));
            if(cur.get(2) == nums2.length - 1)  continue;
            queue.offer(new ArrayList<Integer>(Arrays.asList(cur.get(0), nums2[cur.get(2) + 1], cur.get(2) + 1)));
            
        }
        
        return res;
    }  
}