LC 0632 [H] Smallest Range Covering Elements from K Lists - ALawliet/algorithms GitHub Wiki

class Solution:
    def smallestRange(self, nums: List[List[int]]) -> List[int]:
        # start our heap with the first element from every single list
        heap = [(row[0], i, 0) for i, row in enumerate(nums)]
        heapify(heap)
        
        ans = [-1e9, 1e9]
        
        # temp maximum bound
        right = max(row[0] for row in nums)
        
        while heap:
            left, i, j = heappop(heap)
            # do we need to contract our bound for the final solution?
            if right - left < ans[1] - ans[0]: # smaller range
                ans = left, right
            # if we've finished iterating the current list, so it won't contain the other lists, so rest is invalid and we return the best answer so far
            if j + 1 == len(nums[i]):
                return ans
            # get next element of the list we are currently iterating
            next_num = nums[i][j+1]
            # set new maximum bound
            right = max(right, next_num)
            # put the next element into the heap
            heappush(heap, (next_num, i, j + 1))
            
'''
The key insight/ intuition (non-obvious) is that the heap always contains ONLY 1 element from each list/ row array. The heap NEVER contains 2 or more elements from the SAME list.
This is important because it means that the range calculated at every iteration is always guaranteed to be a range that has exactly 1 element from every single list!
'''

'''
The logic here is that if j + 1 == len(A[i]) that means we explored all the elements of the ith list. Since the output range must include at least one number from each list we can terminate early if we explored all the elements from any one list.
'''