Python's list comprehension: Modify list eleme

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

I had an answer using list comprehension, but @Prune beat me to it. It was really just a cautionary tail, showing how it would be done while making an argument against that approach.

Here's an alternative approach that might fit your needs:

import itertools
import operator

nums = [1,1,0,1,1]
res = itertools.accumulate(nums, operator.and_)

In this case res is an iterable. If you need a list, then

res = list(itertools.accumulate(nums, operator.and_))

Let's break this down. The accumulate() function can be used to generate a running total, or 'accumulated sums'. If only one argument is passed the default function is addition. Here we pass in operator.and_. The operator module exports a set of efficient functions corresponding to the intrinsic operators of Python. When an accumulated and is run on a list of 0's and 1's the result is a list that has 1's up till the first 0 is found, then all 0's after.

Of course we're not limited to using functions defined in the operator module. You can use any function that accepts 2 parameters of the type of the elements in the first parameter (and probably returns the same type). You can get creative, but here I'll keep it simple and just implement and:

import itertools

nums = [1,1,0,1,1]
res = itertools.accumulate(nums, lambda a, b: a and b)

Note: using operator.and_ probably runs faster. Here we're just providing an example using the lambda syntax.

While a list comprehension is not used, to me it has a similar feel. It fits in one line and isn't too hard to read.

Answer 3

Here is a linear-time solution that doesn't mutate global state, doesn't require any other iterators except the nums, and that does what you want, albeit requiring some auxiliary data-structures, and using a seriously hacky list-comprehension:

>>> nums = [1,1,0,1,1]
>>> [f for f, ns in [(1, nums)] for n in ns for f in [f & (n==1)]]
[1, 1, 0, 0, 0]

Don't use this. Use your original for-loop. It is more readable, and almost certainly faster. Don't strive to put everything in a list-comprehension. Strive to make your code simple, readable, and maintainable, which your code already was, and the above code is not.

Answer 4

For a list comprehension approach, you could use index with enumerate:

firstIndex = nums.index(0) if 0 in nums else -1
[1 if i < firstIndex else 0 for i, x in enumerate(nums)]

Another approach using numpy:

import numpy as np
print(np.cumprod(np.array(nums) != 0).tolist())
#[1, 1, 0, 0, 0]

Here we take the convert nums to a numpy array and check to see if the values are not equal to 0. We then take the cumulative product of the array, knowing that once a 0 is found we will multiply by 0 from that point forward.

Answer 5

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)