373 Find K Pairs with Smallest Sums

Posted 2021-01-18 tobeabetterpig

https://leetcode.com/problems/find-k-pairs-with-smallest-sums/discuss/84551/simple-Java-O(KlogK)-solution-with-explanation




public class Solution {
    public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        PriorityQueue<int[]> que = new PriorityQueue<>((a,b)->a[0]+a[1]-b[0]-b[1]);
        List<int[]> res = new ArrayList<>();
        if(nums1.length==0 || nums2.length==0 || k==0) return res;
        for(int i=0; i<nums1.length && i<k; i++) que.offer(new int[]{nums1[i], nums2[0], 0});
        while(k-- > 0 && !que.isEmpty()){
            int[] cur = que.poll();
            res.add(new int[]{cur[0], cur[1]});
            if(cur[2] == nums2.length-1) continue;
            que.offer(new int[]{cur[0],nums2[cur[2]+1], cur[2]+1});
        }
        return res;
    }
}












//  22/27
class Solution {
    public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        
        // generate all pairs and put them into a min heap 
        int[] tmp = new int[2];
        List<int[]> list = new ArrayList<>();
        generate(nums1, nums2, list);
        List<int[]> res = new ArrayList<>();
        
        if(nums1 == null || nums1.length == 0 || nums2 == null || nums2.length == 0) return res;
        PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) -> a[0] + a[1] - b[0] - b[1]);
        for(int[] candidate : list){
            pq.offer(candidate);
        }
        while(k > 0){
            res.add(pq.poll());
            k--;
        }
        return res;
    }
    private void generate(int[] nums1, int[] nums2, List<int[]> list){
        
        for(int num1 : nums1){
            for(int num2 : nums2){
               
                list.add(new int[]{num1, num2}); // list.add(tmp);
                
            }
        }
    }
}