Leetcode solution : find-k-pairs-with-smallest-sums

find_k_pairs_with_smallest_sums.py

# 1) using Min-Heap logN
# 1. make a min heap(Heap Queue : python heapq) S with sum value
# 2. pick first Elem in S and j++ then, push it again
# 3. iterate k times till every j < len(b)
# T(n) = O(klogN)


from heapq import *

class Solution:
    def kSmallestPairs(self, nums1, nums2, k):
        if len(nums1) <= 0 or len(nums2) <= 0:
            return []
        
        h = [(nums1[0] + nums2[0], 0, 0)]
        visited = set()
        res = []
        while len(res) < k and h != []:
            p = heappop(h)
            i1, i2 = p[1], p[2]
            res.append((nums1[i1], nums2[i2]))

            if i1 + 1 < len(nums1) and i1 + 1 not in visited:
                visited.add(i1 + 1)
                newElem = (nums1[i1 + 1] + nums2[0], i1 + 1, 0)
                heappush(h, newElem)
            if i2 + 1 < len(nums2):
                newElem = (nums1[i1] + nums2[i2 + 1], i1, i2 + 1)
                heappush(h, newElem)
        return res
    
# simple test case
if __name__ == "__main__":
    a = Solution()
    print(a.kSmallestPairs([1,7,11], [2,4,6], 3))