Python's list comprehension: Modify list elements if a certain value occurs

Question

How can I do the following in Python's list comprehension?

nums = [1,1,0,1,1]
oFlag = 1
res = []
for x in nums:
    if x == 0:
        oFlag = 0
    res.append(oFlag)
print(res)

# Output: [1,1,0,0,0]

Essentially in this example, zero out the rest of the list once a 0 occurs.

Answer 1

Some context, a list comprehension is a sort of "imperative" syntax for the map and filter functions that exist in many functional programing languages. What you're trying to do is usually referred to as an accumulate, which is a slightly different operation. You can't implement an accumulate in terms of a map and filter except by using side effects. Python allows you have side effects in a list comprehension so it's definitely possible but list comprehensions with side effects are a little wonky. Here's how you could implement this using accumulate:

nums = [1,1,0,1,1]

def accumulator(last, cur):
    return 1 if (last == 1 and cur == 1) else 0

list(accumulate(nums, accumulator))

or in one line:

list(accumulate(nums, lambda last, cur: 1 if (last == 1 and cur == 1) else 0))

Of course there are several ways to do this using an external state and a list comprehension with side effects. Here's an example, it's a bit verbose but very explicit about how state is being manipulated:

class MyState:
    def __init__(self, initial_state):
        self.state = initial_state
    def getNext(self, cur):
        self.state = accumulator(self.state, cur)
        return self.state

mystate = MyState(1)
[mystate.getNext(x) for x in nums]

Answer 2

nums = [1,1,0,1,1]
[int(all(nums[:i+1])) for i in range(len(nums))]

This steps through the list, applying the all operator to the entire sub-list up to that point.

Output:

[1, 1, 0, 0, 0]

Granted, this is O(n^2), but it gets the job done.

Even more effective is simply to find the index of the first 0. Make a new list made of that many 1s, padded with the appropriate quantity of zeros.

if 0 in nums:
    idx = nums.index(0)
    new_list = [1] * idx + [0] * (len(nums) - idx)

... or if the original list can contain elements other than 0 and 1, copy the list that far rather than repeating 1s:

    new_list = nums[:idx] + [0] * (len(nums) - idx)