剑指 Offer II 061. 和最小的 k 个数对

Posted 2021-09-28 炫云云

剑指 Offer II 061. 和最小的 k 个数对

给定两个以升序排列的整数数组 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]
    
输入: 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]

最小堆

与378. 有序矩阵中第 K 小的元素类似，利用最小堆求k小数：

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        if not len(nums1) or not len(nums2): 
            return []
        heap = [(nums1[0] + nums2[0], 0, 0)]
        ans = []
        visited = set()
        while len(heap) and len(ans) < k:
            _, p1, p2 = heapq.heappop(heap) 
            if (p1, p2) in visited: continue
            ans.append((nums1[p1], nums2[p2]))
            visited.add((p1, p2))
            if p1+1 < len(nums1):
                heapq.heappush(heap , (nums1[p1+1]+nums2[p2], p1+1, p2) )
            if p2+1 < len(nums2):
                heapq.heappush(heap , (nums1[p1]+nums2[p2+1], p1, p2+1) )
        return ans

优先级队列（最小堆）

我们可以利用最小堆，在不枚举出全部数对的情况下解决该问题，从而提高效率。

由于nums1和nums2均是已经排好序的升序数组，我们可以固定选择nums1中的一个数字nums[k],让其依次和nums2中的每个数对进行组合从而形成一个以升序排列的数对队列。

例：nums1 = [1,7,11], nums2 = [2,4,6]

这样，我们将原问题转换成在n1个升序队列中，查找最小的前K对数字。为了解决该问题，我们可以参考leetcode 23 合并K个有序链表的方法。我们统计n1个升序队列的队首元素中的最小值，将其加入结果队列，并将指向该队首的指针向后移，直到我们找齐前K对数字。

算法:

维护一个堆，存放n1个升序队列的队首元素，每个元素为一个长度为2数组，里面存放的是nums1和nums2中某个数字的下标。堆中的元素按照数对和由小到大顺序排列，每次我们取出堆中数对和最小的数对，加入结果数组。同时将该数对所在的队列中下一个元素(若存在的话)加入堆中，如此反复直至我们找齐全部K对数字或堆的大小为空。

import heapq
class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        n1 = len(nums1)
        n2 = len(nums2)
        res = []
        queue = []
        for i in range(n1):
            heapq.heappush(queue, (nums1[i] + nums2[0] , i, 0))
        # for i in range(k): #一共弹k次
        while queue and len(res) < k:
            num, x, y = heapq.heappop(queue) #弹出堆里最小一个
            res.append([nums1[x],nums2[y]])
            if y != n2 - 1: #如果这一行还没被弹完
                heapq.heappush(queue , ( nums1[x]+nums2[y+1], x ,y+1 )) #加入num右边的一个值
                
        return res

a = Solution()
print(a.kSmallestPairs( [1,2] ,[3],3))

[[1, 2], [1, 4], [1, 6]]

class Solution:
    def kSmallestPairs(self, nums1, nums2, k):
        queue = []
        def push(i, j):
            if i < len(nums1) and j < len(nums2):
                heapq.heappush(queue, [nums1[i] + nums2[j], i, j])
        push(0, 0)
        pairs = []
        while queue and len(pairs) < k:
            print(len(queue))
            _, i, j = heapq.heappop(queue)

            pairs.append([nums1[i], nums2[j]])
            push(i, j + 1)
            if j == 0: #避免重复遍历
                push(i + 1, 0)
        return pairs

参考

Krahets - 力扣（LeetCode） (leetcode-cn.com)