Python avoid duplicating function call in generator

Viewed 40

Is there a way to write [g(acc, orig) for acc, orig in zip(itertools.accumulate(f(it)), f(it)) in one line without duplicating the call to f? As I understand it, I can't even use := as this is a generator.

1 Answers

You can use itertools.tee(..). This sort of stores the intermediate values while generating (apparently in a single queue). The actual function will be called exactly once as you can see from the number of print statements below.

>>> def fn(x):
...   for i in range(x):
...     print("Yielding", i)
...     yield i

>>> from itertools import tee
>>> seq1, seq2 = tee(fn(10))
>>> [(acc, orig) for acc, orig in zip(accumulate(seq1), seq2)]
Yielding 0
Yielding 1
Yielding 2
Yielding 3
Yielding 4
Yielding 5
Yielding 6
Yielding 7
Yielding 8
Yielding 9
[(0, 0), (1, 1), (3, 2), (6, 3), (10, 4), (15, 5), (21, 6), (28, 7), (36, 8), (45, 9)]
Related