LC 373

Solution to LeetCode 373 Find K Pairs with Smallest Sums.

LeetCode 373

Find K Pairs with Smallest Sums (Medium). [link]

1. Priority queue using min heap. Add the k pairs of first k of nums1 with the first element of nums2 ((0,0), (1,0), … (k,0)) into pq. Every time we pop out the first pair, push the pair with an increased nums2 (e.g. pop (0,0), push (0,1)).

   Time complexity O(k log k). It takes O(k) to select k pairs of numbers. It takes O(logk) time complexity when popping into the heap queue and adjust the order. Space complexity O(k). The pq only stores k pairs of numbers.

class Solution(object):
    def kSmallestPairs(self, nums1, nums2, k):
        """
        :type nums1: List[int]
        :type nums2: List[int]
        :type k: int
        :rtype: List[List[int]]
        """
        result = []
        pq = [(nums1[i] + nums2[0], i, 0) for i in range(min(len(nums1), k))]
        while pq and len(result) < k:
            _, i, j = heappop(pq)
            result.append([nums1[i], nums2[j]])
            if j + 1 < len(nums2):
                heappush(pq, (nums1[i] + nums2[j+1], i, j+1))
        return result

2. Binary search. We aim to find the minimal sum of pairs x such that there are the number of pairs exceed k. We can do binary search with a sliding window, because of the property of this problem:

   • If the sum of pairs is less than x, then the number of pairs is less than k.
   • If the sum of pairs is not less than x, then the number of pairs is no less than k.

   After we get x, we add pairs with sum of them no more than x. Then we add pairs with sum of them equal to x until the total number is k.

   Time complexity O(k + (len(nums1) + len(nums2))* log (range(nums1) + range(nums2))). Binary search takes log (range(nums1) + range(nums2)) time complexity, where range(nums1) = max(nums1) - min(nums1). Sliding window for searching takes O(2 * (k + m + n). Space complexity O(1).

class Solution(object):
    def kSmallestPairs(self, nums1, nums2, k):
        """
        :type nums1: List[int]
        :type nums2: List[int]
        :type k: int
        :rtype: List[List[int]]
        """
        l, r = nums1[0] + nums2[0], nums1[-1] + nums2[-1] + 1
        while l < r: # binary search [l, r)
            m = l + (r - l) // 2
            count = 0
            i, j = 0, len(nums2) -1 
            while i < len(nums1) and j >= 0: # sliding window
                if nums1[i] + nums2[j] > m: # find minimized SUM
                    j -= 1
                else:
                    count += j+1
                    i += 1
            if count < k: # find kth smallest pairs SUM
                l = m + 1
            else:
                r = m
        SUM = l
        result = []
        
        # find pairs with the sum smaller than SUM
        i = len(nums2) -1
        for n1 in nums1:
            while i >= 0 and n1 + nums2[i] >= SUM:
                i -= 1 # find the smallest nums2
            for j in range(i + 1):
                result.append([n1, nums2[j]])
                if len(result) == k:
                    return result
                
        # find pairs with the sum equal to SUM (if len(result) < k)
        i = len(nums2) - 1
        for n1 in nums1:
            while i >= 0 and n1 + nums2[i] > SUM:
                i -= 1
            j = i
            while j >= 0 and n1 + nums2[j] == SUM:
                result.append([n1, nums2[j]])
                if len(result) == k:
                    return result
                j -= 1
        return result