如何获取最后一个元素作为第一个元素的成对迭代器

Question

我正在使用以下函数pairwise来获取有序对的迭代。例如，如果可迭代对象是一个列表[1, 2, 3, 4, 5, 6]，那么我想获取(1, 2), (2, 3), (3, 4), (4, 5), (5, 6), (6, 1). 如果我使用以下函数

from itertools import tee, zip_longest
def pairwise(iterable):
    "s -> (s0,s1), (s1,s2), (s2, s3), ..."
    a, b = tee(iterable)
    next(b, None)
    return zip_longest(a, b)

然后它返回(1, 2), (2, 3), (3, 4), (4, 5), (5, 6), (6, None)。

我在代码中使用dataloader可迭代，因此我只想iterable作为函数的输入传递pairwise，并且不想传递额外的输入。

如何获取第一个元素作为上面提到的最后一项中的最后一个元素？

Answer 1

zip_longest有fillvalue参数

return zip_longest(a, b, fillvalue=iterable[0])

或按照评论中的建议使用next(b, None)in的返回值fillvalue

def pairwise(iterable):
    a, b = tee(iterable)
    return zip_longest(a, b, fillvalue=next(b, None))

输出

print(list(pairwise(lst))) # [(1, 2), (2, 3), (3, 4), (4, 5), (5, 6), (6, 1)]

您也可以在不将列表转换为迭代器的情况下完成此操作

def pairwise(iterable):
    return zip_longest(iterable, iterable[1:], fillvalue=iterable[0])