Compute the cumulative sum of a list until a zero appears

Question

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

list_a = [1, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 1]

I want to get the list_b

Answer 1

Starting Python 3.8, and the introduction of assignment expressions (PEP 572) (:= operator), we can use and increment a variable within a list comprehension:

# items = [1, 1, 1, 0, 1, 1, 0, 1, 0, 1, 1, 1]
total = 0
[total := (total + x if x else x) for x in items]
# [1, 2, 3, 0, 1, 2, 0, 1, 0, 1, 2, 3]

This:

• Initializes a variable total to 0 which symbolizes the running sum
• For each item, this both:
  • either increments total with the current looped item (total := total + x) via an assignment expression or set it back to 0 if the item is 0
  • and at the same time, maps x to the new value of total