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.
'''