Find K paris with smallest sums

用堆就行了,注意k可能大于所有组合数

class Solution(object):
    def kSmallestPairs(self, nums1, nums2, k):
        """
        :type nums1: List[int]
        :type nums2: List[int]
        :type k: int
        :rtype: List[List[int]]
        """
        n, m = len(nums1), len(nums2)
        if not n or not m:
            return []

        result = []
        checker = set()
        heap = [(nums1[0] + nums2[0], 0, 0,)]
        checker.add((0, 0,))
        k = k if k <= n * m else n * m

        while k:
            value, i, j = heapq.heappop(heap)
            result.append([nums1[i], nums2[j]])
            if (i + 1, j,) not in checker and (i + 1 < n):
                checker.add((i + 1, j,))
                heapq.heappush(heap, (nums1[i + 1] + nums2[j], i + 1, j,))
            if (i, j + 1,) not in checker and (j + 1 < m):
                checker.add((i, j + 1,))
                heapq.heappush(heap, (nums1[i] + nums2[j + 1], i, j + 1,))
            k -= 1

        return result

results matching ""

    No results matching ""