Compute cummulative sum until a zero appears

I have a (long) list in which zeros and ones appear at random:

list1 = [1, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 1]

1.I want to get another list

• the sum of the list up to where 0 appears
• where 0 appears, retain 0 in the list

list2 = [3, 0, 2, 0, 1, 0, 3]

It should work for corner cases as well like zeroes at start or end list1 = [0, 1, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 1, 0] list2 = [0, 3, 0, 2, 0, 1, 0, 3, 0]

Answer

You can use itertools.groupby:

from itertools import groupby

list1 = [1, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 1]

out = []
for v, g in groupby(list1):
    if v:
        out.append(sum(g))
    else:
        out.extend(g)  # or out.append(0) if you want to keep only single zero

print(out)

Prints:

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

---

For list1 = [0, 1, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 1, 0] it prints:

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