Leetcode #373. Find K Pairs with Smallest Sums

[yokostan]

Using Priority Queue, beats 17%:

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        PriorityQueue<int[]> queue = new PriorityQueue<>((a, b) -> (a[0] + a[1] - b[0] - b [1]));
        List<List<Integer>> res = new ArrayList<>();
        if (nums1 == null || nums2 == null) return res;

        for (int i = 0; i < nums1.length; i++) {
            for (int j = 0; j < nums2.length; j++) {
                queue.offer(new int[]{nums1[i], nums2[j]});
            }
        }

        while (k > 0 && !queue.isEmpty()) {
            List<Integer> li = new ArrayList<>();
            int[] arr = queue.poll();
            for (int i = 0; i < 2; i++) {
                li.add(arr[i]);
            }
            res.add(li);
            k--;
        }

        return res;
    }
}

PriorityQueue with some tricks, beating 100%:

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<List<Integer>> res = new ArrayList<List<Integer>>();
        if(nums1.length == 0 || nums2.length == 0 || k == 0){
            return res;
        }

        PriorityQueue<int[]> pq = new PriorityQueue<>(k, new Comparator<int[]>(){
            // Overwrite the compare method
            public int compare(int[] a, int[] b) {
                return (a[0] + a[1]) - (b[0] + b[1]);
            }
        });

        for (int i = 0; i < k && i < nums1.length; i++) { //先放起始的K个或者是nums1 length个
            pq.offer(new int[]{nums1[i], nums2[0], 0});// So the array is consist of: the element in nums1, the element in nums2, the curr element's index in nums2.
        }

        // poll the peak element (minElement currently), and offer the next element
        while (k-- > 0 && !pq.isEmpty()) {
            int[] cur = pq.poll();
            res.add(Arrays.asList(cur[0], cur[1]));
            // check the index of nums2 is out of bound or not
            if (cur[2] == nums2.length - 1) {
                continue;
            }
            pq.offer(new int[]{cur[0], nums2[cur[2] + 1], cur[2] + 1});
        }

        return res;

    }
}