[leetcode] 373. Find K Pairs with Smallest Sums

Description

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]

分析

题目的意思是：给定两个数组nums1,nums2，分别从两数组中取出一个数进行配对，求找到k个这样的配对，使得他们的和在所有的配对中最小。

• 解法在暴力的基础上用了一点技巧，把nums1，nums2中可能的组合先列举出来，注意在列举的时候用了min方法，这是因为我们最多不超过k个，所以在单个数组中只取出k个就够了；不够k的话，就把所有的都取出来。然后排个序取前k个配对就行了。

代码

class Solution {
public:
    vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        vector<pair<int,int>> res;
        if(nums1.empty()||nums2.empty()){
            return res;
        }
        int n1=nums1.size();
        int n2=nums2.size();
        for(int i=0;i<min(n1,k);i++){
            for(int j=0;j<min(n2,k);j++){
                res.push_back({nums1[i],nums2[j]});
            }
        }
        sort(res.begin(),res.end(),[](pair<int,int> p1,pair<int,int> p2){
            return p1.first+p1.second<p2.first+p2.second;
        });
        k= k<res.size() ? k:res.size(); 
        return vector<pair<int,int>>(res.begin(),res.begin()+k);
    }
};

分析二

这道题最佳方法是用堆，堆存放的是topK，然后输出堆里面的组合就是最终结果了。python里面是小顶堆，即堆顶的值最小，但是我们要从堆移除较大的数，可以把数值变为负数，这样每次移除的数就是最大的数了。时间复杂度是KKlogK

代码二

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        heap=[]
        for num1 in nums1[:k]:
            for num2 in nums2[:k]:
                if len(heap) <k:
                    heapq.heappush(heap,(-(num1+num2),[num1,num2]))
                else:
                    if(num1+num2<-heap[0][0]):
                        heapq.heappop(heap)
                        heapq.heappush(heap,(-(num1+num2),[num1,num2]))
        return [item[1] for item in heap]

上面的这种方法不是很高效，下面有一个更高效的方法：

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        m=len(nums1)
        n = len(nums2)
        pq = []
        for i in range(len(nums1[:k])):
            pq.append([nums1[i]+nums2[0],i,0])

        ans = []
        while pq and len(ans)<k:
            _, i, j = heappop(pq)
            ans.append([nums1[i], nums2[j]])
            if j+1<n:
                heappush(pq,[nums1[i]+nums2[j+1],i,j+1])
        return ans

参考文献

[LeetCode] Find K Pairs with Smallest Sums 找和最小的K对数字
python中默认是小根堆