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?