核心思路在于如何找到下一个最小的pair加入数组中,故创建一个�heap,结构如下:
- value:这个pair实际代表的数字,以其作为index保证heap的排序正确性
- index:nums1的index
- prime:nums2的index
首先构造初始化的heap,在循环中主要有以下逻辑:
- heapq.heappop(h)
- res.append([nums1[i], nums2[j]])
- 加入新的pair组合
Why Work?
因为在初始化过程中,已经将nums1当中所有index的元素都加入了heap,所以当后续在添加的时候,如何保证其最小性? 我们每次让j代表nums2的index前进一步,这样所有的(i, j+n)当中最小的元素就被我们取出来加入了heap当中,如果存在(i+1, j-n)比当前元素要小,那么由于初始化的关系必然已经存在于当前heap当中,故我们可以保证每次加入新的元素必然会导致最小的元素已经存在于heap当中
Complexity T : O( log(k)k )
class Solution:
def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
if not nums1 or not nums2: return []
h = []
# pair combination
for i in range(len(nums1)):
heapq.heappush(h, (nums1[i] + nums2[0], i, 0))
# return k smallest pairs
res = []
while h and k > 0:
small, i, j = heapq.heappop(h)
res.append([nums1[i], nums2[j]])
if j + 1 < len(nums2):
heapq.heappush(h, (nums1[i] + nums2[j + 1], i, j + 1))
k -= 1
return res