Leetcode 373. Find K Pairs with Smallest Sums

将nums1与nums2组成的所有组合用矩阵M的形式表示，M[i][j]表示nums[i]+nums[j]

用priority queue数据结构进行维护，首先将nums[0] + nums[1]加入pq中，每次pop出一个元素时，将该元素右边和下边未访问过的元素加入pq中。

from Queue import PriorityQueue as PQ
class Solution(object):
    def kSmallestPairs(self, nums1, nums2, k):
        """
        :type nums1: List[int]
        :type nums2: List[int]
        :type k: int
        :rtype: List[List[int]]
        """
        if not nums1 or not nums2:
            return []
        result = []
        # visited记录放入过pq的数
        visited = set()
        pq = PQ()
        # 第一个参数表示sum
        # 第二个参数表示[num1中的index, num2中的index]
        pq.put((nums1[0] + nums2[0], [0,0]))
        visited.add(str([0,0]))
        for i in range(k):
            total, index = pq.get()
            idx1, idx2 = index[0], index[1]
            result.append([nums1[idx1], nums2[idx2]])
            #将右边和下边的元素加入pq中
            neighbors = [[idx1+1, idx2], [idx1, idx2+1]]
            for neighbor in neighbors:
                # 检测是否已经加入过pq
                if str(neighbor) in visited:
                    continue
                # 检测是否越界
                if neighbor[0] >= len(nums1) or neighbor[1] >= len(nums2):
                    continue
                total = nums1[neighbor[0]] + nums2[neighbor[1]]
                pq.put((total, neighbor))
                visited.add(str(neighbor))
            if not pq.queue:
                return result
        return result