Question

我有一个迭代器iterator和一个索引列表indices（可能重复），我想从迭代器中提取这些元素。目前我正在做

indices = sorted(indices)
deltas = [indices[0]] + [indices[i+1] - indices[i] for i in range(len(indices) - 1)]
output = []
for delta in deltas:
    for i in range(delta):
        datum = next(iterator)
    output.append(datum)

这两层循环是否必要？我错过了itertools的技巧吗？

Answer 1

如果内存不是约束，我只会找到最大索引并预先填充迭代器值的数组，直到该最大索引。无论如何，你将不得不计算中间值，所以你真的没有通过计算增量来获得任何东西。

max_index = max(indices)
data = [v for v in itertools.islice(iterator, max_index + 1)]
values = [data[i] for i in indices]

Answer 2

您绝对不需要双循环，因为您可以使用单个循环并且不创建增量但是检查代码变得更复杂：

it = iter(sorted(indices))
index = next(it)
for i, datum in enumerate(iterator):
    if i != index:
        continue
    output.append(datum)
    try:
        index = next(it)
    except StopIteration:
        break

您也可以在列表理解中为非常少量的索引执行此操作，因为您需要支票的开销（但避免使用sort）：

[datum for i, datum in enumerate(x) if i in indices]

您可以将indices转换为set来降低支票费用。我有兴趣看到sort在set构造上的表现（集查找为O（1））：

indices = set(indices)
[datum for i, datum in enumerate(x) if i in indices]

更新：第一和第三个选项在时间上大致相当于超过900毫秒（从第一个边缘开始），从10,000,000个项目中选择1000个随机索引。 OP的代码大约运行了1.2秒。

索引到Python迭代器

2 个答案: