Problem statement

https://leetcode.com/problems/smallest-range-covering-elements-from-k-lists/

Solution

Let n be total number of elements in all nums and m be number of lists in nums. Then brute-force solution is to check all possible O(n^2) possible ranges and then check in O(n) if we have every element in this range, so complexity will be O(n^3). We can improve it to O(n^2 log n), if we use binary search.

There is much better approach. Let us keep heap with one candidate from each list and at each step move one candidate to the right (we can look at this process as sliding window among all numbers from all lists). What candidate we need to move? The one, which is smallest, because if we move any other element, window size can only increase. So, each moment we find smallest element in our heap, pop it and add new element - next element from corresponding list. What we keep in list is tuple (number, number of list, index in list). We also update our ans and cur_max.

Complexity

Time complexity is O(n log m), space complexity is O(m).

Code

class Solution:
    def smallestRange(self, nums):
        heap = [(i[0], j, 0) for j, i in enumerate(nums)]
        cur_max = max(heap)[0]
        heapify(heap)
        ans = (-float("inf"), float("inf"))
        
        while True:
            ans = min(ans, (heap[0][0], cur_max), key = lambda x: x[1] - x[0])
            num, num_list, ind = heappop(heap)
            if ind == len(nums[num_list]) - 1: break
            heappush(heap, (nums[num_list][ind+1], num_list,ind+1) )
            cur_max = max(cur_max, nums[num_list][ind+1])
            
        return ans