Question

Formatted question description: https://leetcode.ca/all/373.html

 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]

 @tag-heap

Algorithm

The value of num1[0]+num2[0] must be the smallest, so we first throw this value into the heap.

Then start traversing the next element.

The next value must be the smallest one from num1[1]+num2[0] and num1[0]+num2[1].

Loop all the combinations in turn until the extracted element is K.

Code

Java

import java.util.*;

public class Find_K_Pairs_with_Smallest_Sums {

    public static void main(String[] args) {
        Find_K_Pairs_with_Smallest_Sums out = new Find_K_Pairs_with_Smallest_Sums();
        Solution s = out.new Solution();

        // output: [[1,2],[1,4],[1,6]]
//        s.kSmallestPairs(new int[]{1,7,11}, new int[]{2,4,6}, 3).forEach(e -> System.out.println(e.toString()));

        s.kSmallestPairs(new int[]{1,1,2}, new int[]{1,2,3}, 2).forEach(e -> System.out.println(e.toString()));
    }

    public class Solution_NlogN {
        public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {

            List<List<Integer>> result = new ArrayList<>();

            if(nums1.length == 0 || nums2.length == 0 || k == 0) {
                return result;
            }

            PriorityQueue<int[]> q = new PriorityQueue<>(
                (a, b) -> a[0]+a[1]-b[0]-b[1]
            ); // @note: min at top of heap

            // all nums-1 to in heap
            for(int i = 0; i < nums1.length && i < k; i++) {

                // last 0 meaning index of nums, to check if this index is hitting nums end
                int[] oneVector = new int[]{nums1[i], nums2[0], 0};
                q.offer(oneVector);
            }

            while(k > 0 && !q.isEmpty()){

                k--;

                int[] pair = q.poll();
                result.add(Arrays.asList(pair[0], pair[1]));

                if(pair[2] == nums2.length-1) {
                    continue;
                }

                int nextIndexInNums1 = pair[0];
                int nextIndexInNums2 = pair[2] + 1;
                q.offer(new int[]{nextIndexInNums1, nums2[nextIndexInNums2], nextIndexInNums2});
            }

            return result;
        }
    }

    class Solution {
        public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {

            List<List<Integer>> result = new ArrayList<>();

            if(nums1.length == 0 || nums2.length == 0){
                return result;
            }

            PriorityQueue<List<Integer>> queue = new PriorityQueue<>(
                k, // initial size will not block auto-grow
                (a,b) -> b.get(0) + b.get(1) - a.get(0) - a.get(1)
            );

            for (int i = 0; i < Math.min(k, nums1.length); i++) {
                for (int j = 0; j < Math.min(k, nums2.length); j++) {
                    if (queue.size() < k) {
                        queue.offer(Arrays.asList(nums1[i], nums2[j]));
                    } else if (nums1[i] + nums2[j] < queue.peek().get(0) + queue.peek().get(1)) {
                        queue.offer(Arrays.asList(nums1[i], nums2[j]));
                        queue.poll();
                    }
                }
            }

            return new ArrayList<>(queue);
        }
    }
}

All Problems

All Solutions