Leetcode Find K Pairs with smallest sums

时间：2023-03-08 17:11:54

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]

本题的特点在于两个list nums1和nums2都是已经排序好的。本题如果把所有的(i, j)组合都排序出来，再取其中最小的K个。其实靠后的很多组合根本用不到，所以效率较低，会导致算法超时。为了简便，把nums1和nums2里面的元素写成A1,A2,...,A5, B1,...,B5. 维护一个最小堆。并逐渐往里面push and pop元素。

注意以下几点：1. 如果(Ai, Bj)大于堆顶的元素，那么没有必要把（Ai, Bj）和它之后的元素Push进堆。2. 为了便于管理，解法里其实只维护了一个指针idx2。如果 （A1,B[idx2]）已经小于堆顶的元素，那么就把(A1,B[idx2]),...,(A5,B[idx2])全都推到堆里面。

第〇步，把一个最大元素放进堆。注意这一步很重要，否则可能会出现(A1,B2)大于所有(Ai,B1)的情况，也就是第二步中红色箭头大于蓝色箭头的情况始终不会出现，那么如果k足够大，就可能会出现pop一个空堆的bug。

第一步 把所有的(A1，B1),..., (A5,B1)推入堆。也就是把所有的黑色箭头连接的pair都放到堆里。

Leetcode Find K Pairs with smallest sums

然后在逐个pop堆顶的元素。直到出现(A1，B2)小于堆顶的情况 （红色箭头小于蓝色箭头）。如果出现，把所有的(A1,B2),...,(A5,B2)放进堆。恢复pop堆顶元素，并且开始监视下一个红色箭头(A1,B3)是否小于堆顶的元素。

Leetcode Find K Pairs with smallest sums

 def kSmallestPairs(nums1, nums2, k):

     """

     :type nums1: List[int]

     :type nums2: List[int]

     :type k: int

     :rtype: List[List[int]]

     """

     ans = []

     heap = [(0x7FFFFFFF, None, None)]

     size1, size2 = len(nums1), len(nums2)

     idx2 = 0

     while len(ans) < min(k, size1 * size2):

         if idx2 < size2:

             sum, i, j = heap[0]

             if nums2[idx2] + nums1[0] < sum:

                 for idx1 in range(size1):

                     heapq.heappush(heap, (nums1[idx1] + nums2[idx2], idx1, idx2))

                 idx2 += 1

         sum, i, j = heapq.heappop(heap)

         ans.append((nums1[i], nums2[j]))

     return ans

相关文章