Increment list based on number pattern

Question

I have a list of zeros and ones that looks like this:

lst = [0, 1, 1, 1, 1, 0, 0, 0, 1, 1, 0, 0, 0, 1, 0, 1, 1]

How can I transform this lst into this:

transformed_lst = lst = [0, 1, 1, 1, 1, 0, 0, 0, 2, 2, 0, 0, 0, 3, 0, 4, 4]

Basically, at each occurance of a 1, transform it to an n+1 integer. I'm sure there is an elegant way to do this with itertools/groupby/functools. Here is an attempt, but not quite correct:

from itertools import cycle

ints = cycle(range(len(lst))) 
transformed_lst = [next(ints) if i != 0 in lst else 0 for i in lst]  

>>> [0, 0, 1, 2, 3, 0, 0, 0, 4, 5, 0, 0, 0, 6, 0, 7, 8]  

Answer 1

You basically have two states - "reading 0s" and "reading 1s" - and when you switch between then (namely from ones to zeroes) the delta to be applied for subsequent 1s change:

reading_zeroes = True
delta = 0
for x in input:
    if x:
        reading_zeroes = False
        x += delta
    elif not reading_zeroes:
        delta += 1
        reading_zeroes = True
    yield x

Answer 2

using itertools.count(),itertools.chain() and itertools.groupby():

In [14]: from itertools import *

In [15]: c=count(1)

In [16]: lis=[0, 1, 1, 1, 1, 0, 0, 0, 1, 1, 0, 0, 0, 1, 0, 1, 1]

In [17]: list(chain(*[list(g) if k!=1 else [next(c)]*len(list(g)) for k,g in groupby(lis)]))

Out[17]: [0, 1, 1, 1, 1, 0, 0, 0, 2, 2, 0, 0, 0, 3, 0, 4, 4]

here you can also use sum(1 for _ in g) in place of len(list(g))

As Demanded, a readable version using generator function:

In [27]: def func(l):
    c=count(1)
    for k,g in groupby(l):
        if k==1:
            for x in [next(c)]*sum(1 for _ in g):
                yield x
        else:
            for x in g:
                yield x
   ....:                 

In [28]: list(func(lis))
Out[28]: [0, 1, 1, 1, 1, 0, 0, 0, 2, 2, 0, 0, 0, 3, 0, 4, 4]