632. Smallest Range Covering Elements from K Lists

This stupid problem… I’ve read the solution 2 months ago and failed to understand it. I know it will keep hunting me and today it comes back in another problem list…

Idea

In this section, don’t worry too much about how we implement this using code, just get the rough idea, in the end, the idea of how we solve it matters, the code will naturally comes out when the idea is clear.

I had a hard time to find the solution by bare hands without considering how to solve it by writing code, until I understand this:

To cover all the lists, we have to take one element from each list, and since the lists are sorted, we can start from the first element of each list, and somehow move forward.

The starting point is here:

1

Take note:

  • the max of them is 5, min of them is 0
  • range [0,5] obviously covers all 3 lists (think a bit, we took the min and max…)

We want to minimize the range, so we have two options:

  • reduce the max value, this is not possible, simply because the max value(5 in this case) is the min value of the list it belongs to(in this case the 3rd list), if we reduce it, we lost the row.
  • increase the min value, we have to do this.

Next step is to increase the min value to the next value in the same list:

2

We see the min become 4 and max is 9 now, this range is not smaller than the previous one ([0,5]) so we ignore it.

Continue increasing the min element (1st row, 4 for now).

Until this point:

3

Only until this case, we have [20,24] which has a smaller range than all the previous ranges. We save this info somewhere(so later we can return it)

When to stop? Kind of obvious, when we cannot do the above anymore, which is when we exhaust one of the list!

Now we can return the smallest range.

Algorithm

Hint on the data structures we need to use:

  • we need to find smallest element -> heap
  • then how do we track the max? simply keep a cmax variable… so we get min element from heap and max element from the var

I know that my solution below is quite verbose comparing to the official answer or those posted in the solution section, but I believe verbose code is easier to understand both for the candidate and the interviewer.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
INF=float('inf')
class Solution:
    def smallestRange(self, nums: List[List[int]]) -> List[int]:
        # heap, pointers to each list's current smallest element
        pts=[]
        cmax=-INF
        # collect smallest element in each list, and save its row index and column index
        for i in range(len(nums)):
            num=nums[i][0]
            pts.append((num,i,0))
            cmax=max(cmax,num)
        
        heapq.heapify(pts)
        # keep track of the best range we have found
        res_start=pts[0][0]
        res_end=cmax
        res_range=res_end-res_start
        while True:
            num,row,idx=heapq.heappop(pts)
            # exhaused, we're done
            if idx+1>=len(nums[row]):
                break
            next_num=nums[row][idx+1]
            heapq.heappush(pts,(next_num,row,idx+1))
            cmax=max(cmax,next_num)
            next_range=cmax-pts[0][0]
            # update only when current range is better
            if next_range<res_range:
                res_start=pts[0][0]
                res_end=cmax
                res_range=res_end-res_start

        return [res_start,res_end]