373. Find K Pairs with Smallest Sums

https://leetcode.com/problems/find-k-pairs-with-smallest-sums/

Use heap & BFS

需要用一个dict来做visited记录,否则会有重复

import heapq
class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        queue = []
        res = []
        visited = dict()
        if len(nums1) == 0 or len(nums2) == 0:
            return res
        heapq.heappush(queue,(nums1[0]+nums2[0],0,0))
        visited[(0,0)] = 1
        while len(res) < k and queue:
            curr = heapq.heappop(queue)
            currSum,i,j = curr
            res.append([nums1[i],nums2[j]])
            if i+1 < len(nums1) and (i+1,j) not in visited:
                heapq.heappush(queue,(nums1[i+1]+nums2[j],i+1,j))
                visited[(i+1,j)] = 1
            if j+1 < len(nums2) and (i,j+1) not in visited:
                heapq.heappush(queue,(nums1[i]+nums2[j+1],i,j+1))
                visited[(i,j+1)] = 1
        return res

Last updated

Was this helpful?