373 查找和最小的K对数字（多路归并、小根堆）

1. 问题描述：

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

示例 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]

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/find-k-pairs-with-smallest-sums

2. 思路分析：

① 一开始没有思路的时候可以画画图，画图是一个能够很好分析题目模型和规律的一个方法，nums1中的每一个数字可以与nums2中的每一个数字组合，由题可知，nums1和nums2都是升序的数组，所以可以得到n个每一个长度为m的升序序列如下图所示，我们需要求解的是这n个序列中前k对和最小的数字，很明显这是一个经典的多路归并问题。对于多路归并问题都是类似的套路。首先我们需要将n个序列中的第一个数字加入到小根堆中，使用堆是为了维护k对和最小的数对（凡是求解前k小或者k大的问题都可以考虑使用堆进行维护），而且因为我们需要知道当前第t个序列数字的下一个数字，所以需要将当前序列数字组合对应的nums1和nums2的下标加入到堆中，所以在堆中需要维护三个信息，第一个是当前序列中的来自nums1和nums2的数字组合对应和，当前数字组合中对应的num1数字下标，nums2数字下标，因为使用的是python语言，所以使用heapq模块将列表queue维护成一个小根堆，将上面需要保存的三个信息作为元组的成员（元组可以很方便存储多个信息）加入到堆中。

② 具体的做法（多路归并都是类似的套路）：一开始初始化堆，将n个序列中的第一个数字也即Ai + B0，nums1的下标，nums2的下标作为python中元组的成员加入到小根堆中。当k大于0并且堆不为空的时候弹出堆顶元素，堆顶元素存储了当前和最小的数字，将当前弹出的堆顶元素的最小值加入到答案中，并且将当前数字的下一个数字，对应的nums1中数字的下标，nums2中数字的下标作为元素成员加入到堆中即可。

3. 代码如下：

from typing import List
import heapq


class Solution:
    # 多路归并思想
    def kSmallestPairs(self, a: List[int], b: List[int], k: int) -> List[List[int]]:
        # python中的heapq模块可以将列表维护为一个堆
        queue = list()
        n, m = len(a), len(b)
        res = list()
        for i in range(n):
            # 根据AB的组合将n个序列的第一个元素加入到堆中, 将元素封装为元组会比较方便
            heapq.heappush(queue, (a[i] + b[0], i, 0))
        while k > 0 and len(queue):
            # 弹出堆顶元素
            poll = heapq.heappop(queue)
            # 当前堆顶元素肯定是和最小的组合
            res.append([a[poll[1]], b[poll[2]]])
            # 根据画出的图可以很快确定对应的下标, 只要是当前的nums2的数字还没到最终那么就将对应的数字加入到堆中
            if poll[2] + 1 < m:
                heapq.heappush(queue, (a[poll[1]] + b[poll[2] + 1], poll[1], poll[2] + 1))
            k -= 1
        return res