leetcode373. 查找和最小的K对数字

1.题目描述

给定两个以升序排列的整数数组 nums1 和 nums2 , 以及一个整数 k 。

定义一对值 (u,v)，其中第一个元素来自 nums1，第二个元素来自 nums2 。

请找到和最小的 k 个数对 (u1,v1),  (u2,v2)  ...  (uk,vk) 。

示例 1:

输入: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
输出: [1,2],[1,4],[1,6]
解释: 返回序列中的前 3 对数：
     [1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]
示例 2:

输入: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
输出: [1,1],[1,1]
解释: 返回序列中的前 2 对数：
     [1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]
示例 3:

输入: nums1 = [1,2], nums2 = [3], k = 3 
输出: [1,3],[2,3]
解释: 也可能序列中所有的数对都被返回:[1,3],[2,3]
 

提示:

1 <= nums1.length, nums2.length <= 104
-109 <= nums1[i], nums2[i] <= 109
nums1, nums2 均为升序排列
1 <= k <= 1000

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/find-k-pairs-with-smallest-sums
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

2.解题思路

解法一：

最小堆

https://leetcode-cn.com/problems/find-k-pairs-with-smallest-sums/solution/133-cha-zhao-he-zui-xiao-de-kdui-shu-zi-jv2j6/

解法二：

最大堆

https://leetcode-cn.com/problems/find-k-pairs-with-smallest-sums/solution/python-wei-hu-zui-da-dui-by-shuai-dao-me-sdpf/

3.代码实现

方式一：

class Solution(object):
    def kSmallestPairs(self, nums1, nums2, k):
        """
        :type nums1: List[int]
        :type nums2: List[int]
        :type k: int
        :rtype: List[List[int]]
        """
        # https://leetcode-cn.com/problems/find-k-pairs-with-smallest-sums/solution/133-cha-zhao-he-zui-xiao-de-kdui-shu-zi-jv2j6/
        l1 = len(nums1)
        l2 = len(nums2)
        i,j = 0,0
        res = []
        queue = []
        def push(i,j):
            if i < l1 and j < l2:
                heapq.heappush(queue, (nums1[i] + nums2[j] , i, j))
        push(0,0)
        while queue and len(res) < k:
            _, i ,j = heapq.heappop(queue)
            res.append([nums1[i],nums2[j]])
            push(i, j + 1)
            # 防止重复加入
            if j == 0:
                push(i+1,j)
        
        return res

# 相当于以下代码
        l1 = len(nums1)
        l2 = len(nums2)
        visited = set()
        i,j = 0,0
        res = []
        queue = []
        def push(i,j):
            if i < l1 and j < l2:
                heapq.heappush(queue, (nums1[i] + nums2[j] , i, j))
                visited.add((i, j))
        push(0,0)
        visited.add((0,0))
        while queue and len(res) < k:
            _, i ,j = heapq.heappop(queue)
            res.append([nums1[i],nums2[j]])
            if (i,j+1) not in visited:
                push(i, j + 1)
            if (i+1,j) not in visited:
                push(i+1,j)
        
        return res

方式二：

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        # https://leetcode-cn.com/problems/find-k-pairs-with-smallest-sums/solution/python-wei-hu-zui-da-dui-by-shuai-dao-me-sdpf/
        heap = []
        for n1 in nums1:
            for n2 in nums2:
                if len(heap) < k:
                    heapq.heappush(heap, (-n1-n2, [n1, n2]))
                elif heap and -heap[0][0] > n1 + n2:
                    heapq.heappop(heap)
                    heapq.heappush(heap, (-n1-n2, [n1, n2]))
                else:
                    break
        return [heapq.heappop(heap)[1] for _ in range(k) if heap][::-1]
        # 会报错
        # return [heapq.heappop(heap)[1] for _ in  heap]