leetcode 373. Find K Pairs with Smallest Sums

发表于 | 分类于

You are given two integer arrays nums1 and nums2 sorted in ascending order and an integer k.

Define a pair (u,v) which consists of one element from the first array and one element from the second array.

Find the k pairs (u1,v1),(u2,v2) …(uk,vk) with the smallest sums.

Example 1:


Given nums1 = [1,7,11], nums2 = [2,4,6],  k = 3

Return: [1,2],[1,4],[1,6]

The first 3 pairs are returned from the sequence:
[1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]

Example 2:


Given nums1 = [1,1,2], nums2 = [1,2,3],  k = 2

Return: [1,1],[1,1]

The first 2 pairs are returned from the sequence:
[1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]

Example 3: 


Given nums1 = [1,2], nums2 = [3],  k = 3

Return: [1,3],[2,3]

All possible pairs are returned from the sequence:
[1,3],[2,3]

  • mysolution the limit exceeded 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
class Solution(object):
    def kSmallestPairs(self, nums1, nums2, k):
        """
        :type nums1: List[int]
        :type nums2: List[int]
        :type k: int
        :rtype: List[List[int]]
        """
        tmp = []
        if not nums1 or not nums2:
            return []
        for i in nums1:
            for j in nums2:
                tmp.append([i+j,i,j])
        tmp = sorted(tmp,key=lambda x:x[0])
        out = []
        for l in tmp:
            out.append(l[1:])
        if k > len(out):return out
        return out[:k]
  • a o(k*logk) solution
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
from heapq import *
class Solution:
    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 []
        
        heap = [(nums1[0] + nums2[0], 0, 0)] 
        visited = set()
        output = []
        
        while len(output) < k and heap:
            val = heappop(heap)
            output.append([nums1[val[1]], nums2[val[2]]])
            
            if val[1] < len(nums1) - 1 and (val[1]+1, val[2]) not in visited:
                visited.add((val[1]+1, val[2]))
                heappush(heap, (nums1[val[1] + 1] + nums2[val[2]], val[1] + 1, val[2]))
            if val[2] < len(nums2) - 1 and (val[1], val[2] + 1) not in visited:
                visited.add((val[1], val[2] + 1))
                heappush(heap, (nums1[val[1]] + nums2[val[2] + 1], val[1], val[2] + 1))
        
        return output