Search code examples
pythoniteratorpython-itertools

How to get pairwise iterator with last element as being the first

I am using the following function pairwise to get the iteration of ordered pairs. For example, if the iterable is a list [1, 2, 3, 4, 5, 6] then I want to get (1, 2), (2, 3), (3, 4), (4, 5), (5, 6), (6, 1). If I use the following function

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)

then it returns (1, 2), (2, 3), (3, 4), (4, 5), (5, 6), (6, None).

I am using dataloader in my code as iterable, so I want to pass only iterable as input to the function pairwise and I don't want to pass extra inputs.

How do I get the first element as the last element in the last item as mentioned above?

Solution

  • zip_longest has fillvalue parameter

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

    or as suggested in the comments use the returned value of the next(b, None) in fillvalue

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

    Output

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

    You can also do it without converting the list to iterators

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