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?
Advertisement
Answer
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])