def sort(lists, k=300):
    n = len(lists)
    ans = []
    heap = [(lists[i][0], i) for i in range(n)]
    idx = {i: 1 for i in range(n)}
    heapq.heapify(heap)
    for _ in range(k):
        val, i = heapq.heappop(heap)
        ans.append(val)
        heapq.heappush(heap, (lists[i][idx[i]], i))
        idx[i] += 1
    
    return ans

大概是这样