跳转至

LeetCode: 373. 查找和最小的K对数字

1、题目描述

给定两个以升序排列的整形数组 nums1nums2, 以及一个整数 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]

2、解题思路

​ 借助堆来实现

  • 首先用用第一个数组的第0个元素,加上第二个数组的所有的值,放入堆中
  • 然后开始从堆里里面取值,如果取出的值大于第一个数字第1个元素加上第二个数组第0个元素的值,这时候,我们将第一个数组的第1各元素,与第二个数组的所有元素匹配,放入堆中

实例说明

[1 1 2] [1 2 3]
c = 0
一开始,将第一个元数组的第0个元素,与第二个数组所有的数字匹配,放入堆中
(2,0,0)
(3,0,1)
(4,0,2)
c=1
将上面的组合放到堆中,接着,我们开始判断,从堆中取出来的组合,(2,0,0)的2,是不是大于第一个数组中下一个元素,与第二个数组中第0各元素之和,如果大于,结果中存放这个数字,并且将刚刚取出来的放入堆中,并且将第二个元素与剩余所有元素的匹配放入堆中
class Solution:
    def kSmallestPairs(self, nums1, nums2, k):
        """
        :type nums1: List[int]
        :type nums2: List[int]
        :type k: int
        :rtype: List[List[int]]
        """

        if not nums1 or not nums2:
            return []

        length1 = len(nums1)
        length2 = len(nums2)

        count = min(length1 * length2, k)

        result = []

        heap = [(nums1[0] + nums2[i], 0, i) for i in range(length2)]
        heapq.heapify(heap)
        c = 1

        for i in range(count):
            if not heap:

                for j in range(length2):
                    heapq.heappush(heap, (nums1[c] + nums2[j], c, j))
                c += 1
            temp = heapq.heappop(heap)

            if c < length1:
                temp_sum = nums1[c] + nums2[0]
                if temp[0] < temp_sum:
                    result.append([nums1[temp[1]], nums2[temp[2]]])
                else:
                    result.append([nums1[c], nums2[0]])
                    heapq.heappush(heap, temp)
                    for j in range(1, length2):
                        heapq.heappush(heap, (nums1[c] + nums2[j], c, j))
                    c += 1
            else:
                result.append([nums1[temp[1]], nums2[temp[2]]])

        return result