Using itertools.tee
to zip an iterable with itself
Since you are iterating twice on the same iterable, one possibility is to use itertools.tee
:
from itertools import tee
def zip_offset(iterable, offset):
a, b = tee(iterable)
for _ in range(offset):
next(b, None)
return zip(a, b)
print(list( zip_offset(range(10), 3) ))
# [(0, 3), (1, 4), (2, 5), (3, 6), (4, 7), (5, 8), (6, 9)]
Using tee
is particularly important if the iterable is an iterator. Consider:
from itertools import tee, islice
v = map(lambda x: 2*x, [0,1,2,3,4,5,6,7,8,9])
a, b = tee(v)
print(list( zip(a, islice(b, 3, None)) ))
# [(0, 6), (2, 8), (4, 10), (6, 12), (8, 14), (10, 16), (12, 18)]
v = map(lambda x: 2*x, [0,1,2,3,4,5,6,7,8,9])
print(list( zip(v, islice(v, 3, None)) ))
# [(0, 8), (10, 12), (14, 16)]
more_itertools.zip_offset
Alternatives from a module are welcome as well
Meet module more_itertools
and its function zip_offset
:
from more_itertools import zip_offset
print(list( zip_offset('01234', 'abcdefghijkl', offsets=(0, 3)) ))
# [('0', 'd'), ('1', 'e'), ('2', 'f'), ('3', 'g'), ('4', 'h')]
You can look at its source code:
from itertools import islice
from itertools import chain, repeat # used when negative offsets
from itertools import zip_longest # used when longest=True
def zip_offset(*iterables, offsets, longest=False, fillvalue=None):
if len(iterables) != len(offsets):
raise ValueError("Number of iterables and offsets didn't match")
staggered = []
for it, n in zip(iterables, offsets):
if n < 0:
staggered.append(chain(repeat(fillvalue, -n), it))
elif n > 0:
staggered.append(islice(it, n, None))
else:
staggered.append(it)
if longest:
return zip_longest(*staggered, fillvalue=fillvalue)
return zip(*staggered)
In the special case when the offset is 1, and assuming you're using python>=3.10, there is also pairwise
in itertools
:
from itertools import pairwise
print(list( pairwise('abcdef') ))
# [('a', 'b'), ('b', 'c'), ('c', 'd'), ('d', 'e'), ('e', 'f')]
print(f"\n{x=}, {y=}")
with modern f-strings. \$\endgroup\$print(f"\nx={x}, y={y}")
I assume \$\endgroup\${x=}
is a shorthand forx={x}
, they do the same thing. \$\endgroup\$