根据另一个列表返回连续出现的列表

Question

I have a list with the following number:

[0, 0, 0, 0, 1, 1, 0, 0, 1]

What I wanna do is count how many consecutive 1s I have and generate a new list like the following one:

[0, 0, 0, 0, 1, 2, 0, 0, 1]

Is there a way to do it in a more compact way than this method?

lst = [0, 0, 0, 0, 1, 1, 0, 0, 1]
new_lst = []
counter = 0
for i in range(len(lst)):
    if (lst[i] == 1): counter+=1
    else:counter = 0
    new_lst.append(counter)
print(new_lst)

Answer 1

You could take advantage of the walrus operator if you're using python 3.8 or later:

[counter := counter + 1 if v else 0 for v in lst]

or perhaps itertools.accumulate :

[v for v in itertools.accumulate(lst, lambda acc, v:acc + v if v else 0)]

In both case the result is:

[0, 0, 0, 0, 1, 2, 0, 0, 1]

Answer 2

It's not so much compact than the question solution but it's more clear:

l = [0, 0, 0, 0, 1, 1, 0, 0, 1]
s = []

occ = 0
for i in l:
    occ = occ+1 if i==1 else 0
    s.append(occ)
print(s)

Output:

[0, 0, 0, 0, 1, 2, 0, 0, 1]