[LeetCode] 373. Find K Pairs with Smallest Sums 找和最小的K对数字

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:

Given nums1 = [1,7,11], nums2 = [2,4,6],  k = 3

Return: [1,2],[1,4],[1,6]

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:

Given nums1 = [1,1,2], nums2 = [1,2,3],  k = 2

Return: [1,1],[1,1]

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:

Given nums1 = [1,2], nums2 = [3],  k = 3 

Return: [1,3],[2,3]

All possible pairs are returned from the sequence:
[1,3],[2,3] 

Credits:
Special thanks to @elmirap and @StefanPochmann for adding this problem and creating all test cases.

给定2个以升序排列的整数数组和1个整数k，从2个数组中各拿出一个数字组成一对，找出k对和最小的数字组合。

解法1：最简单的想法就是暴力brute force解法，但效率肯定不高。

解法2: 最小堆。把所有的点对加入到最小堆，然后输出前k个。但没有利用到“两个数组都有序”这个条件，就算数组无序，也可以利用这个方法。要利用有序这个条件，可以借助mergesort的思路，pair的第一个元素至多包含了nums1数组的前k个元素，k以后的可以不用考虑。所以，这形成了k个list，每一个list都包含了nums2的元素。每一次取所有list中的最小值，然后该list下一个元素入队。

Java：

 public List<int[]> kSmallestPairs(int[] nums1, int[] nums2, int k) {
    	List<int[]> r = new ArrayList<>();
    	if(nums1.length == 0 || nums2.length == 0) return r;
    	int size = Math.min(nums1.length, k);
    	int[] index = new int[size];
    	PriorityQueue<int[]> queue = new PriorityQueue<>(new Comparator<int[]>(){
			@Override
			public int compare(int[] o1, int[] o2) {
				Integer s1 = o1[0] + o1[1];
				Integer s2 = o2[0] + o2[1];
				return s1.compareTo(s2);
			}
    	});
    	for(int i = 0; i < size; i++){
    		queue.add(new int[]{nums1[i], nums2[0], i});
    	}
    	int count = 0;
    	while(!queue.isEmpty()){
    		int[] pair = queue.poll();
    		r.add(new int[]{pair[0], pair[1]});
    		int id = pair[2];
    		if(++index[id] < nums2.length)
    			queue.add(new int[]{nums1[id], nums2[index[id]], id});
    		count++;
    		if(count == k)
    			break;
    	}
    	return r;
    }　　

Python:

from heapq import heappush, heappop

class Solution(object):
    def kSmallestPairs(self, nums1, nums2, k):
        """
        :type nums1: List[int]
        :type nums2: List[int]
        :type k: int
        :rtype: List[List[int]]
        """
        pairs = []
        if len(nums1) > len(nums2):
            tmp = self.kSmallestPairs(nums2, nums1, k)
            for pair in tmp:
                pairs.append([pair[1], pair[0]])
            return pairs

        min_heap = []
        def push(i, j):
            if i < len(nums1) and j < len(nums2):
                heappush(min_heap, [nums1[i] + nums2[j], i, j])

        push(0, 0)
        while min_heap and len(pairs) < k:
            _, i, j = heappop(min_heap)
            pairs.append([nums1[i], nums2[j]])
            push(i, j + 1)
            if j == 0:
                push(i + 1, 0)  # at most queue min(n, m) space
        return pairs

C++:

// Time:  O(k * log(min(n, m, k))), where n is the size of num1, and m is the size of num2.
// Space: O(min(n, m, k))

class Solution {
public:
    vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        vector<pair<int, int>> pairs;
        if (nums1.size() > nums2.size()) {
            vector<pair<int, int>> tmp = kSmallestPairs(nums2, nums1, k);
            for (const auto& pair : tmp) {
                pairs.emplace_back(pair.second, pair.first);
            }
            return pairs;
        }

        using P = pair<int, pair<int, int>>;
        priority_queue<P, vector<P>, greater<P>> q;
        auto push = [&nums1, &nums2, &q](int i, int j) {
            if (i < nums1.size() && j < nums2.size()) {
                q.emplace(nums1[i] + nums2[j], make_pair(i, j));
            }
        };

        push(0, 0);
        while (!q.empty() && pairs.size() < k) {
            auto tmp = q.top(); q.pop();
            int i, j;
            tie(i, j) = tmp.second;
            pairs.emplace_back(nums1[i], nums2[j]);
            push(i, j + 1);
            if (j == 0) {
                push(i + 1, 0);  // at most queue min(m, n) space.
            }
        }
        return pairs; 
    }
};

　　

　　

 

All LeetCode Questions List 题目汇总