Leetcode Find K Pairs with smallest sums

Posted 2020-08-24 lettuan

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 (u₁,v₁),(u₂,v₂) ...(u_k,v_k) 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都放到堆里。

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

 

 1 def kSmallestPairs(nums1, nums2, k):
 2     """
 3     :type nums1: List[int]
 4     :type nums2: List[int]
 5     :type k: int
 6     :rtype: List[List[int]]
 7     """
 8     ans = []
 9     heap = [(0x7FFFFFFF, None, None)]
10     size1, size2 = len(nums1), len(nums2)
11     idx2 = 0
12     while len(ans) < min(k, size1 * size2):
13         if idx2 < size2:
14             sum, i, j = heap[0]
15             if nums2[idx2] + nums1[0] < sum:
16                 for idx1 in range(size1):
17                     heapq.heappush(heap, (nums1[idx1] + nums2[idx2], idx1, idx2))
18                 idx2 += 1
19         sum, i, j = heapq.heappop(heap)
20         ans.append((nums1[i], nums2[j]))
21     return ans