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

  • 0

    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) for itr in iters]
            for _ in xrange(k):
                if not heap: break
                (_, num1, num2), itr = heapq.heappop(heap)
                ans.append([num1, num2])
                    heapq.heappush(heap, (, itr))
                except StopIteration:
            return ans

Log in to reply

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