Min range to cover all lists

minrange.py

#!/usr/bin/env python

"""
Given $k$ sorted lists, each of size $n$, find a range $(p,q)$ such that:
    1. q-p is minimized
    2. for any list, there is an element $r$, such that $p <= r <= q$

Example:
    [4, 10, 15, 24]
    [0, 9, 12, 20]
    [5, 18, 22, 30]
    result: (20, 24) with 24-20 = 4

Source of problem: https://www.youtube.com/watch?v=zplklOy7ENo
"""

def minrange(sorted_lists):
    k = len(sorted_lists)
    n = len(sorted_lists[0])
    assert all(len(lst) for lst in sorted_lists), "some list are empty"
    assert all(len(lst) == n for lst in sorted_lists), "uneven list lengths"
    assert all(all(x <= y for x,y in zip(lst, lst[1:])) for lst in sorted_lists), "not all lists are sorted"
    idx = [0] * k
    # generate initial solution: assume
    p = min(sorted_lists[i][idx[i]] for i in range(k))
    q = max(sorted_lists[i][idx[i]] for i in range(k))
    # iterate: O(k) for min/max, O(nk) iterations for while loop
    while True:
        # find the min to move up
        i = min(range(k), key = lambda i: sorted_lists[i][idx[i]])
        if idx[i] + 1 >= n:
            break # no more
        else:
            idx[i] += 1
        new_p = min(sorted_lists[i][idx[i]] for i in range(k))
        new_q = max(sorted_lists[i][idx[i]] for i in range(k))
        if new_q - new_p < q - p:
            p, q = new_p, new_q
    return p, q

def main():
    sorted_lists = [
        [4, 10, 15, 24],
        [0, 9, 12, 20],
        [5, 18, 22, 30],
    ]
    print(minrange(sorted_lists))

if __name__ == "__main__":
    main()