leetcode 373. 查找和最小的K对数字 - 不带visited数组的表格图遍历 - 字典嵌套初始化

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

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

思路

用位置关系和堆来取前k个最小数对

方案1 基础代码 超时

visited数组控制重复
每次添加目标元素下边和右边的元素入堆

超时可能的原因：
visited数组初始化让时间复杂度变成了O(n平方)

取前k个最小数对的步骤里面，时间复杂度是O(klog2(n))

import heapq
class Solution:
    def kSmallestPairs(self, nums1, nums2, k):
        visited = [[False]*len(nums2) for _ in range(len(nums1))]
        candi = []
        ans = []
        i = j = 0
        heapq.heappush(candi,(nums1[i]+nums2[j],i,j))
        visited[i][j] = True
        count = 0
        while len(candi)>0 and k>0:
            count +=1
            s, i, j = heapq.heappop(candi)
            ans.append([nums1[i],nums2[j]])
            k-=1
            if i <len(nums1)-1 and not visited[i+1][j]:
                heapq.heappush(candi,(nums1[i+1]+nums2[j],i+1,j))
                visited[i+1][j] = True
            if j<len(nums2)-1 and not visited[i][j+1]:
                heapq.heappush(candi,(nums1[i]+nums2[j+1],i,j+1))
                visited[i][j+1] = True
        return ans

方案1 优化1 通过

某个元素加入堆的时候，只考虑[i+1][j-1],[i-1][j+1] 是否被访问过

import heapq
class Solution:
    def kSmallestPairs(self, nums1, nums2, k):
        ans = []
        candi = []
        visited = collections.defaultdict(lambda :collections.defaultdict(lambda :False))
        for i in range(len(nums1)+1):
            visited[i][-1]=True
        for j in range(len(nums2)+1):
            visited[-1][j]=True
        i = j = 0
        heapq.heappush(candi, (nums1[i] + nums2[j], i, j))
        while len(candi) and k>0:
            s, i, j = heapq.heappop(candi)
            ans.append([nums1[i],nums2[j]])
            # ans.append([i,j])
            visited[i][j] = True
            k -= 1
            if i<len(nums1)-1 and visited[i+1][j-1]:
                heapq.heappush(candi,(nums1[i+1]+nums2[j],i+1,j))
            if j<len(nums2)-1 and visited[i-1][j+1]:
                heapq.heappush(candi,(nums1[i]+nums2[j+1],i,j+1))
        return ans

方案1 优化2 通过

在上一个优化的基础上，不必考虑所有边界，优化此段代码

某个位置被访问到再进行赋值

import heapq
class Solution:
    def kSmallestPairs(self, nums1, nums2, k):
        ans = []
        candi = []
        visited = collections.defaultdict(lambda :collections.defaultdict(lambda :False))
        i = j = 0
        heapq.heappush(candi, (nums1[i] + nums2[j], i, j))
        while len(candi) and k>0:
            s, i, j = heapq.heappop(candi)
            ans.append([nums1[i],nums2[j]])
            # ans.append([i,j])
            visited[i][j] = True
            k -= 1
            if j==0 : visited[i + 1][j - 1] = True
            if i==0 : visited[i - 1][j + 1] = True
            if i<len(nums1)-1 and visited[i+1][j-1]:
                heapq.heappush(candi,(nums1[i+1]+nums2[j],i+1,j))
            if j<len(nums2)-1 and visited[i-1][j+1]:
                heapq.heappush(candi,(nums1[i]+nums2[j+1],i,j+1))
        return ans

方案2 其他人的题解

与方案1的思路相近，但是控制重复的方法不太一样

遇到 j==0 的时候才加入下一行的数据，否则直接加入右列的数据

class Solution:
    def kSmallestPairs(self, nums1, nums2, k):
        h = []
        from heapq import heappop,heappush
        def push(i, j):
            if i < len(nums1) and j < len(nums2):
                heappush(h, [nums1[i] + nums2[j], i, j])
        push(0, 0)
        pairs = []
        while h and len(pairs) < k:
            _, i, j = heappop(h)
            pairs.append([nums1[i], nums2[j]])
            push(i, j + 1)
            if j == 0: #避免重复遍历
                push(i + 1, 0)
        return pairs

收获

1. 边界visited数组使用时赋值
2. 矩阵访问控制去重的方法
3. 字典的嵌套初始化
   visited = collections.defaultdict(lambda :collections.defaultdict(lambda :False))