如何高效地在迭代器上进行“n-wise”迭代

Question

可能是重复的问题，但我找不到任何相关内容。

我有一个非常长的迭代器（10000个元素），我需要每次迭代500个元素。如果我的迭代器是range(10000)，那么代码应该如下所示：

Iteration #1: 0, 1, 2, ... 497, 498, 499
Iteration #2: 1, 2, 3, ... 498, 499, 500
Iteration #3: 2, 3, 4, ... 499, 500, 501
Iteration #4: 3, 4, 5, ... 500, 501, 502
...
Iteration #9500: 9499, 9500, 9501 ... 9996, 9997, 9998
Iteration #9501: 9500, 9501, 9502 ... 9997, 9998, 9999

等等。这里有一个方法：

def nwise_slice(lst, n):
    for i in range(len(lst) - n + 1):
        yield lst[i:i + n]

然而，这种方法不能与惰性迭代器一起使用。我尝试使用迭代器创建解决方案，并从 itertools 的 pairwise 和 consume 配方中进行了适应（请参见此处）以创建以下内容：

import itertools

def nwise_iter(lst, n):
    iters = itertools.tee(lst, n)
    for idx, itr in enumerate(iters):
        next(itertools.islice(itr, idx, idx), None)

    for group in zip(*iters):
        yield group

这个方法与使用list相同（虽然返回的是一个tuple而不是list，但对我来说无关紧要）。我认为它不会创建很多不必要的切片。这种解决方案适用于非可切片迭代器，如文件（我计划使用）。然而，itertools的解决方案速度慢了2倍：

In [4]: %timeit list(nwise_slice(list(range(10000)), 500))
46.9 ms ± 729 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [5]: %timeit list(nwise_iter(list(range(10000)), 500))
102 ms ± 3.95 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

我不想把所有的测试数据加载到内存中，以利用slice方法。有没有更有效的方法来完成这个操作？

- iz_

---

请查看我在此答案中的grouper()函数。 - martineau

@martineau 谢谢。不幸的是，它迭代了(1, 2, 3), (4, 5, 6)而不是(1, 2, 3), (2, 3, 4), (3, 4, 5), (4, 5, 6)。 - iz_

你需要在迭代器前进时保持迭代器结果有效吗？你需要对每个子序列进行O(1)索引访问吗？ - Davis Herring

@DavisHerring 不需要那些。我可以承担“搞砸”原始迭代器的风险。 - iz_

你可以使用滑动窗口算法。虽然不一定更快，但更简洁的方法是查看 more_itertools 库 list(more_itertools.windowed(range(10000), 500))。 - pylang

Answer 1

使用双端队列(deque)来实现"记忆化(memoize)"你的项，这个怎么样？

from collections import deque

def nwise_slice(it, n):
    deq = deque((), n)
    for x in it:
        deq.append(x)
        if len(deq)==n: yield deq

my_range = range(8)
for sub in nwise_slice(my_range, 5):
    print(sub)
# =>
# deque([0, 1, 2, 3, 4], maxlen=5)
# deque([1, 2, 3, 4, 5], maxlen=5)
# deque([2, 3, 4, 5, 6], maxlen=5)
# deque([3, 4, 5, 6, 7], maxlen=5)

- Walter Tross

---

谢谢你的回答。幸运的是，它非常快。但是，我得到了错误的结果 - 它每次都产生相同的deque。（另外，lst[0:n]将在迭代器上失败，但我认为这可以通过使用islice来修复？） - iz_

我想主要问题在于我对Python还是有些初学者。无论如何，我添加了我的测试代码，也许你能够找出我是否使用不当，或者我的用例与你的有何不同之处。 - Walter Tross

有趣的是：迭代它没问题，但调用 list(nwise_slice(my_range, 5)) 就会出错。嗯... - iz_

@Tomothy32：使用list是无用的，因为每个结果都是相同的对象deq。这就是为什么我问过关于使之前的结果失效的问题... - Davis Herring

@WalterTross 如果没有它，list(nwise_slice(x, 5)) 将会是相同对象的列表。这是意料之外的。可以像使用 itertools.groupby 一样线性使用结果，但是和 itertools.groupby 一样，你必须知道生成器不返回副本并自己制作一个 [(k, list(g)) for k,g in groupby(...)] 作为对象或迭代器在下一次迭代中将不再有效。而不是复制，可以将返回结果包装在可失效代理中，在控制返回到生成器时变为无效。 - Dan D.