[LeetCode] 373. Find K Pairs with Smallest Sums

 Problem: https://leetcode.com/problems/find-k-pairs-with-smallest-sums/

Assume len(nums1) = M and len(nums2) = N.

nums1 and nums2 can form M sorted list and each list contains N elements.

This problem equals to merge M sorted list and output the first k elements.

The first element of each list is ( nums1[i],  nums2[0] ).  

Because nums1 is sorted list,  ( nums1[i],  nums2[0] ) is the smallest sum on each list.

To find the smallest sum among all list (global smallest sum), we put the smallest sum on each list into the heap.

Then popup the global smallest sum.   Then expend the heap with the next smallest sum of the list from where the current global smallest sum comes.

Time Complexity = O ( K * Log N ).  N = len(nums1)

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        """
                1 ,    2 ,    3
            1  (1,1)  (1,2)  (1,3)
            
            1  (1,1)  (1,2)  (1,3)
            
            2  (2,1)  (2,2)  (2,3)
            
        
        """
        
        if k == 0 or not nums1 or not nums2: return []
        
        heap = []
        result = []
        
        # add the first element of all lines into heap
        for y in range(len(nums1)):
            heapq.heappush(heap, (nums1[y] + nums2[0], y, 0))
        
        while len(result) < k and heap:
            # pop current minimum sum from heap
            _, y, x = heapq.heappop(heap)
            result.append([nums1[y], nums2[x]])
            
            # push next sum combination into heap
            x += 1
            if x < len(nums2):
                heapq.heappush(heap, (nums1[y] + nums2[x], y, x))
        
        return result