I'm looking for a nice way to sequentially combine two itertools operators. As an example, suppose we want to select numbers from a generator sequence greater than a threshold, after having gotten past that threshold. For a threshold of 12000, these would correspond to it.takewhile(lambda x: x<12000)
and it.takewhile(lambda x: x>=12000)
:
# Set up an example generator:
def lcg(a=75,c=74,m=2**16+1,x0 = 1):
xn = x0
yield xn
while True:
xn = (a*xn+c) % m
yield xn
# First 20 elements:
list(it.islice(lcg(), 20))
[1, # <- start sequence, start it.takewhile(lambda x: x<12000)
149,
11249, # <- last element of it.takewhile(lambda x: x<12000)
57305, # <- start it.takewhile(lambda x: x>=12000) here
38044,
35283,
24819,
26463,
18689,
25472, # <- last element of it.takewhile(lambda x: x>=12000); end of sequence
9901,
21742,
57836,
12332,
7456,
34978,
1944,
14800,
61482,
23634]
Is there a way to select the sequence of greater than 12000, including the initial values less than 12000, i.e. the desired output is:
[1, 149, 11249, 57305, 38044, 35283, 24819, 26463, 18689, 25472]
This is trivial to do with two for-loops, but I'm looking for an itertools-type way (maybe a one-liner?) of combining the two operators without reseting the lcg
generator.
One approach to writing a one-liner for the problem with the existing itertools
library would be to use a flag variable with {0}
as a default value to indicate which predicate to use. At first the flag evaluates to a truthy value so that the first predicate (x < 12000
) is made effective, and if the first predicate fails, pop the set so the flag becomes falsey to make the second predicate (x >= 12000
) effective. By popping 0
from the set, it also allows the expression to fall back to the second predicate on the same iteration when the first predicate fails:
takewhile(lambda x, f={0}: f and (x < 12000 or f.pop()) or x >= 12000, lcg())
Note that it's safe to use a mutable object as the default value of an argument in this case because lambda
creates a new function object along with a new instance of the mutable object every time it's evaluated.