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

user_2700041 0
Viewed 37

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?

itertools iterator python
2 Answers

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])
user_51680110

You can simply "append" the first element of the iterable to the iterable first.

def pairwise(iterable):
    first = next(iterable)
    new_iterable = chain([first], iterable, [first])

    a, b = tee(new_iterable)
    next(b, None)
    return zip_longest(a, b)
user_11268410
Related