Sequential chaining of itertools operators

Question:

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 less 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.

Asked By: njp

||

Answers:

One approach to writing a one-liner for the problem with the existing itertools library would be to use a mutable flag variable such as a set to indicate which predicate to use. Make its default value truthy so that the first predicate (x < 12000) is made effective, and if the first predicate fails, set the flag to falsey to use the second predicate (x >= 12000) for evaluation:

takewhile(lambda x, f={1}: 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.

Demo: https://replit.com/@blhsing/VirtualGuiltyRatio

Answered By: blhsing