Search code examples
pythonpython-itertools

itertools.takewhile() accessing next element in python


I faced the problem that I dont understand how itertools.takewhile() code works.

import itertools

z = [3,3,9,4,1]

zcycle = itertools.cycle(z)
next_symbol = zcycle.next()
y = list(itertools.takewhile(lambda symbol: symbol == next_symbol or symbol == 9, zcycle))
print y

my code suppose to give me elements of the list from the beginning if they are the same or if element is equal to 9. So once we hit element that differs from previous one we should stop.

I expected that the outcome would be [3, 3] but instead I got [3, 9]. Why we miss the very first element of the list? and is it possible somehow to get output equal to [3, 3, 9]?


Solution

  • You removed the first 3 from the sequence here:

    next_symbol = zcycle.next()
    

    That advances zcycle to the next element, so it'll yield one more 3, not two.

    Don't call next() on the zcycle object; perhaps use z[0] instead:

    next_symbol = z[0]
    

    Now zcycle will yield 3, then another 3, then 9, after which the takewhile() condition will be False and all iteration will stop.