632. Smallest Range Covering Elements from K Lists

Problem Link

Description

---

You have k lists of sorted integers in non-decreasing order. Find the smallest range that includes at least one number from each of the k lists.

We define the range [a, b] is smaller than range [c, d] if b - a < d - c or a < c if b - a == d - c.

 

Example 1:

Input: nums = [[4,10,15,24,26],[0,9,12,20],[5,18,22,30]]
Output: [20,24]
Explanation: 
List 1: [4, 10, 15, 24,26], 24 is in range [20,24].
List 2: [0, 9, 12, 20], 20 is in range [20,24].
List 3: [5, 18, 22, 30], 22 is in range [20,24].

Example 2:

Input: nums = [[1,2,3],[1,2,3],[1,2,3]]
Output: [1,1]

 

Constraints:

• nums.length == k
• 1 <= k <= 3500
• 1 <= nums[i].length <= 50
• -105 <= nums[i][j] <= 105
• nums[i] is sorted in non-decreasing order.

Solution

---

Python3

class Solution:
    def smallestRange(self, nums: List[List[int]]) -> List[int]:
        N = len(nums)
        res = [-inf, inf]
        pq = []
        curr = []
        count = [0] * N
        covered = 0
        currMax = -inf
 
        for i, arr in enumerate(nums):
            heappush(pq, (arr[0], 0, i))
        
        while pq:
            x, arrIndex, numsIndex = heappop(pq)
 
            if count[numsIndex] == 0:
                covered += 1
            count[numsIndex] += 1
            heappush(curr, (x, arrIndex, numsIndex))
            currMax = max(currMax, x)
 
            while covered == N:
                a, b = curr[0][0], currMax
                c, d = res
 
                if b - a < d - c or (b - a == d - c and a < c):
                    res = [a, b]
 
                _, arrIndex2, numsIndex2 = heappop(curr)
                count[numsIndex2] -= 1
                if count[numsIndex2] == 0:
                    covered -= 1
                if arrIndex2 + 1 < len(nums[numsIndex2]):
                    heappush(pq, (nums[numsIndex2][arrIndex2 + 1], arrIndex2 + 1, numsIndex2))
 
        return res