Pythonic O(klogn) solution: heap, iterator, and generator


  • 0
    M

    The idea is to put the sum-iterator pairs into the min-heap.

    import heapq
    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 and nums2):
                return []
    
            def generator(num1, nums2):
                for num2 in nums2:
                    yield (num1 + num2, num1, num2)
    
            ans = []
            iters = [generator(num1, nums2) for num1 in nums1]
            heap = [(itr.next(), itr) for itr in iters]
    
            for _ in xrange(k):
                if not heap: break
                (_, num1, num2), itr = heapq.heappop(heap)
                ans.append([num1, num2])
                try:
                    heapq.heappush(heap, (itr.next(), itr))
                except StopIteration:
                    pass
            return ans
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.