632. Smallest Range Covering Elements from K Lists

Problem

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.

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

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
-10⁵ <= nums[i][j] <= 10⁵
nums[i] is sorted in non-decreasing order.

Test Cases

1 2	class Solution: def smallestRange(self, nums: List[List[int]]) -> List[int]:

solution_test.py

import pytest

from solution import Solution


@pytest.mark.parametrize('nums, expected', [
    ([[4,10,15,24,26],[0,9,12,20],[5,18,22,30]], [20,24]),
    ([[1,2,3],[1,2,3],[1,2,3]], [1,1]),

    (
        [
            [55, 197, 202, 272, 456, 514, 529, 838, 888, 890, 918, 996],
            [133, 272, 387, 431, 522, 729, 770, 817, 900, 902],
            [35, 100, 175, 324, 416, 439, 658, 824, 886, 989],
            [15, 57, 108, 127, 189, 313, 405, 489, 499, 516, 608, 628, 687, 742, 853, 925],
            [111, 132, 209, 244, 264, 291, 317, 380, 470, 481, 490, 635, 688, 889, 946, 986],
            [76, 83, 332, 427, 596, 612, 660, 673, 764, 844, 964, 978, 997, 999],
            [5, 18, 146, 244, 344, 402, 420, 470, 549, 686, 697, 708, 768, 778, 802, 876],
            [58, 153, 290, 352, 386, 535, 638, 777, 826, 829, 852, 894, 937, 947],
            [190, 192, 367, 426, 453, 469, 563, 594, 654, 699, 792, 828, 857, 942, 960, 992, 1000],
            [3, 94, 115, 175, 393, 400, 423, 465, 521, 526, 674, 844, 872, 885, 919],
        ],
        [844, 900]
    ),
    (
        [
            [5, 7, 10, 24, 29, 32, 34, 35, 35, 43, 53, 55, 55, 58, 69, 89, 92, 94, 100],
            [0, 1, 4, 6, 7, 8, 14, 29, 31, 42, 45, 56, 73, 74, 75, 79, 79, 86, 92, 100],
            [0, 2, 25, 28, 33, 37, 44, 51, 58, 69, 72, 73, 77, 81, 84, 100],
            [18, 43, 50, 56, 65, 67, 69, 79, 96, 96],
            [1, 2, 5, 12, 31, 31, 54, 60, 76, 79],
            [3, 14, 28, 32, 55, 56, 65, 68, 81, 93, 95],
            [2, 7, 22, 24, 27, 30, 33, 40, 46, 69, 70, 72, 78, 84, 94, 97, 97, 97, 100],
            [1, 5, 7, 9, 10, 17, 25, 35, 38, 39, 41, 50, 57, 79, 84, 93, 94],
            [3, 6, 7, 9, 15, 20, 24, 26, 28, 33, 38, 43, 47, 54, 58, 62, 77, 79, 81, 96],
            [6, 6, 19, 21, 37, 38, 42, 44, 59, 62, 70, 80, 91, 96, 100],
            [0, 3, 7, 7, 8, 14, 15, 37, 40, 75, 78, 84, 93],
            [10, 21, 26, 31, 41, 64, 68, 77, 78, 100],
            [4, 15, 17, 27, 28, 30, 32, 35, 35, 36, 37, 44, 48, 76, 88],
        ],
        [68, 79]
    ),
])
class Test:
    def test_solution(self, nums, expected):
        sol = Solution()
        assert sol.smallestRange(nums) == expected

Thoughts

跟 23. Merge k Sorted Lists 差不多。

同样拿 k 个数组的第一个数字构建大小为 k 的最小堆。堆顶数值就是当前区间的左端点。右端点可以在这 k 个数字进堆的时候记录下来最大值得到。

每次用堆顶对应数组的下一个数字替换堆顶，并恢复堆，此时堆的最小值和最大值就是新区间的边界。最小值直接取自堆顶，最大值则在原本最大值和刚替换进来的数字之间取较大的。

如果堆顶对应数组元素都用完了则停止。

时间复杂度 O(k * n * log k)，其中 n 是 nums 的平均长度。空间复杂度 O(k)。

Code

solution.py

class Solution:
    def smallestRange(self, nums: list[list[int]]) -> list[int]:
        k = len(nums)
        leaf = k >> 1
        heap = [(arr[0], arr, 0) for arr in nums] # min heap

        def shift_down(pos: int):
            while pos < leaf:
                left = (pos << 1) + 1
                right = left + 1
                child = right if right < k and heap[right][0] < heap[left][0] else left
                if heap[pos][0] > heap[child][0]:
                    heap[pos], heap[child] = heap[child], heap[pos]
                    pos = child
                else:
                    return

        # Build the heap.
        for i in range(leaf - 1, -1, -1):
            shift_down(i)

        min_l = heap[0][0]
        r = max(arr[0] for arr in nums)
        min_len = r - min_l
        while (j := heap[0][2]) < len(heap[0][1]) - 1:
            heap[0] = (heap[0][1][j+1], heap[0][1], j+1)
            r = max(r, heap[0][0])
            shift_down(0)
            if (cur_len := r - heap[0][0]) < min_len:
                min_len = cur_len
                min_l = heap[0][0]

        return [min_l, min_l + min_len]