复杂的heapq延迟合并(最小化使用的空间)

Question

给定一个排序列表列表，我希望生成一个排序输出列表。

这很容易：

nums : List[List[int]]
h = heapq.merge(nums)

然而，我也希望用它产生的内部列表的索引来标记输出的每个元素。例如

nums = [[1,3,5], [2,6], [4]]
h = _ # ???
for x in h:
print(x)
# Outputs
# (1,0)
# (2,1)
# (3,0)
# (4,2)
# (5,0)
# (6,1)

我写了一个有效的版本，

h = heapq.merge(*map(lambda l: map(lambda x: (x,l[0]), l[1]), enumerate(nums)))

但我担心我可能已经失去了一个理想的空间复杂性保证；我如何才能知道(转换的(内心清单是否正在显现？(在我的尝试中，*究竟做了什么？(

Answer 1

由于默认的元组比较是字典式的，因此可以将最里面的元素x转换为(x, i)，其中i是nums中包含x的列表的索引。例如，

import heapq
from itertools import repeat
nums = [[1,3,5], [2,6], [4]]
nums_with_inds = (zip(lst, repeat(i)) for i, lst in enumerate(nums))
res = heapq.merge(*nums_with_inds)
for tup in res:
print(tup)
# (1, 0)
# (2, 1)
# (3, 0)
# (4, 2)
# (5, 0)
# (6, 1)