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 大概是这样