373. Find K Pairs with Smallest Sums - Medium

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]

 

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

因为两个数组都是有序的，往min heap中添加元素的理想顺序应该是：首先只考虑两个数组的前k个元素，对应nums1中的某个元素，下一个理想顺序的组合应该是 nums1中的这个元素 + 从堆中poll出来的对应的nums2元素的下一个元素。堆中存index比较方便，先把(i, 0)放入堆中，i对应nums1的index，0是nums2的index；poll出来一个最小值之后，保存结果，检查index是否越界，然后把下一个candidate放入堆中

basic idea: use minHeap to keep track on next minimum pair sum, and we only need to maintain K possible candidates in the data structure.

for every numbers in nums1, its best candidate(yields min sum) always starts from nums2[0] since arrays are all sorted; and for a specific number in nums1, its next candidate sould be [this specific number] + nums2[current_associated_index + 1], unless out of boundary

时间：O(KlogK)，空间：O(K)

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

 

之前的方法用了min heap，但类似于brute force，能过但是太慢了，O(K^2 logK)

class Solution {
    public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<int[]> res = new ArrayList<>();
        if(nums1 == null || nums2 == null) return res;
        PriorityQueue<int[]> minHeap = new PriorityQueue<>((a, b) -> (b[0] + b[1]) - (a[0] + a[1]));
        int m = Math.min(nums1.length, k), n = Math.min(nums2.length, k);
        for(int i = 0; i < m; i++) {
            for(int j = 0; j < n; j++) {
                minHeap.offer(new int[] {nums1[i], nums2[j]});
                if(minHeap.size() > k)
                    minHeap.poll();
            }
        }
        while(!minHeap.isEmpty()) {
            res.add(0, minHeap.poll());
        }
        return res;
    }
}

or

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 || k == 0) {
            return res;
        }
        PriorityQueue<int[]> minHeap = new PriorityQueue<>((a, b) -> (nums1[a[0]] + nums2[a[1]]) - (nums1[b[0]] + nums2[b[1]]));
        boolean[][] visited = new boolean[nums1.length][nums2.length];
        minHeap.offer(new int[] {0, 0});
        visited[0][0] = true;
        
        for(int i = 0; i < k && !minHeap.isEmpty(); i++) {
            int[] cur = minHeap.poll();
            int r = cur[0], c = cur[1];
            
            List<Integer> tmp = new ArrayList<>();
            tmp.add(nums1[r]);
            tmp.add(nums2[c]);
            res.add(tmp);
            
            if(r + 1 < nums1.length && !visited[r + 1][c]) {
                minHeap.offer(new int[] {r + 1, c});
                visited[r + 1][c] = true;
            }
            if(c + 1 < nums2.length && !visited[r][c + 1]) {
                minHeap.offer(new int[] {r, c + 1});
                visited[r][c + 1] = true;
            }
        }
        
        return res;
    }
}