如何修改pairwise迭代器，让最后元素与首个元素配对（仅传iterable参数）

要让pairwise函数生成循环的相邻有序对（最后一对为最后一个元素和第一个元素），可以通过tee复制多个迭代器，并将第二个迭代器的末尾拼接第一个元素来实现，无需额外参数。修改后的代码如下：

from itertools import tee, chain

def pairwise(iterable):
    # 复制出三个独立的迭代器，分别用于完整遍历、偏移遍历、获取首元素
    full_iter, offset_iter, first_iter = tee(iterable, 3)
    
    try:
        # 获取第一个元素，用于最后拼接
        first_elem = next(first_iter)
    except StopIteration:
        # 处理空可迭代对象的情况
        return iter([])
    
    # 让偏移迭代器前进一位，从第二个元素开始
    next(offset_iter, None)
    
    # 将偏移迭代器与包含首元素的迭代器链接，实现循环
    offset_iter = chain(offset_iter, [first_elem])
    
    # 拼接完整迭代器和处理后的偏移迭代器
    return zip(full_iter, offset_iter)

代码说明 #

1. tee复制迭代器：通过tee(iterable, 3)创建三个独立的迭代器，避免直接重复迭代原对象（尤其适合DataLoader这类无法重复迭代的对象）。
2. 获取首元素：通过first_iter获取原可迭代对象的第一个元素，用于最后拼接成循环对。
3. 偏移迭代器：让offset_iter前进一位，确保它从第二个元素开始遍历。
4. 循环拼接：用chain将offset_iter的末尾拼接首元素，使得偏移迭代器的最后一个元素是原对象的第一个元素。
5. 生成有序对：通过zip将完整迭代器和处理后的偏移迭代器配对，得到目标结果。

测试验证 #

输入列表[1, 2, 3, 4, 5, 6]，调用pairwise后会返回：
(1, 2), (2, 3), (3, 4), (4, 5), (5, 6), (6, 1)

该实现兼容PyTorch DataLoader、生成器等各类可迭代对象，无需传入额外参数。

内容的提问来源于stack exchange，提问作者hanugm