Find K Pairs with Smallest Sums

1package pkg5leetcode.official75;2 3/*4 * Find K Pairs with Smallest Sums | LC 3735 * APPROACH: Min-heap seed (0,j); expand next column pairs.6 * COMPLEXITY: Time O(k log k), Space O(k)7 */8import java.util.*;9 10public class official75_LC373FindKPairsWithSmallestSums {11    static List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {12        List<List<Integer>> res = new ArrayList<>();13        if (nums1.length == 0 || nums2.length == 0 || k == 0) return res;14        PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) ->15                Integer.compare(nums1[a[0]] + nums2[a[1]], nums1[b[0]] + nums2[b[1]]));16        for (int j = 0; j < Math.min(nums2.length, k); j++) pq.offer(new int[]{0, j});17        while (!pq.isEmpty() && res.size() < k) {18            int[] cur = pq.poll();19            res.add(Arrays.asList(nums1[cur[0]], nums2[cur[1]]));20            if (cur[0] + 1 < nums1.length) pq.offer(new int[]{cur[0] + 1, cur[1]});21        }22        return res;23    }24 25    public static void main(String[] args) {26        List<List<Integer>> r = kSmallestPairs(new int[]{1,7,11}, new int[]{2,4,6}, 3);27        check(r.equals(Arrays.asList(Arrays.asList(1,2), Arrays.asList(1,4), Arrays.asList(1,6))), "case1");28        List<List<Integer>> r2 = kSmallestPairs(new int[]{1,1,2}, new int[]{1,2,3}, 2);29        check(r2.equals(Arrays.asList(Arrays.asList(1,1), Arrays.asList(1,1))), "case2");30        System.out.println("all tests passed");31    }32 33    static void check(boolean cond, String name) {34        if (!cond) throw new AssertionError("FAILED: " + name);35        System.out.println("  PASS " + name);36    }37}