[LeetCode 373] 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 (u1,v1),(u2,v2) ...(uk,vk) with the smallest sums.

Example 1:

Input: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
Output: [[1,2],[1,4],[1,6]] 
Explanation: 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:

Input: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
Output: [1,1],[1,1]
Explanation: 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:

Input: nums1 = [1,2], nums2 = [3], k = 3
Output: [1,3],[2,3]
Explanation: All possible pairs are returned from the sequence: [1,3],[2,3]

Key observation: K pairs && smallest. -> Heap

Min Heap solution. O(K * log K) runtime, O(K) space

1. create min heap that takes a pair's index and compares pairs' sum.

2. add the first Math.min(nums1.length, k) pairs in format (i, 0) to min heap. i is nums1's index, 0 represents 0 index of nums2.

3. poll from min heap, add the pair to answer, then create a new pair with nums2 index incremented by 1 and add it back to min heap.

4. repeat 3 until either there are k pairs or there is no pairs left in min heap.

A variation of this algorithm is to only add the smallest pair(0, 0) to min heap. Then each time we poll a pair, we try to increase the pair's 1st index; then try to increase the pair's 2nd index. However, we'll need some extra book keeping to avoid adding the same pair more than once.

class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<List<Integer>> res = new ArrayList<>();
        if(nums1.length == 0 || nums2.length == 0) {
            return res;
        }
        PriorityQueue<int[]> minPq = new PriorityQueue<>((a, b) -> {
            return nums1[a[0]] + nums2[a[1]] - nums1[b[0]] - nums2[b[1]];
        });
        Set<Integer>[] set = new Set[nums1.length];
        for(int i = 0; i < nums1.length; i++) {
            set[i] = new HashSet<>();
        }
        minPq.add(new int[]{0, 0});
        set[0].add(0);
        while(k > 0 && minPq.size() > 0) {
            int[] curr = minPq.poll();
            List<Integer> list = new ArrayList<>();
            list.add(nums1[curr[0]]);
            list.add(nums2[curr[1]]);
            res.add(list);
            if(curr[0] < nums1.length - 1 && !set[curr[0] + 1].contains(curr[1])) {
                minPq.add(new int[]{curr[0] + 1, curr[1]});
                set[curr[0] + 1].add(curr[1]);
            }
            if(curr[1] < nums2.length - 1 && !set[curr[0]].contains(curr[1] + 1)) {
                minPq.add(new int[]{curr[0], curr[1] + 1});
                set[curr[0]].add(curr[1] + 1);
            }
            k--;
        }
        return res;
    }
}

[LeetCode 373] Find K Pairs with Smallest Sums

猜你喜欢