Search code examples
pythonlist-comprehensionfor-comprehension

python list comprehension not working


I have this very simple problem: I need to set the values in list a to 1 for each index in list b:

>>> a=[0, 0, 0, 0]
>>> b=[1, 3]

the desired result then would be:

[0, 1, 0, 1]

The elegant solution, if python was worth its salt, would be this of course:

>>> a[b]=1
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: list indices must be integers, not list

But of course that doesn't work... so I've tried the following list comprehensions, but they both produce the same error, as if the comprehension was causing it (on the for):

First the simple version I was really expecting to work:

>>> a[x for x in b] = 1
  File "<stdin>", line 1
    a[x for x in b] = 1
          ^
SyntaxError: invalid syntax

Then the slightly more complex version:

>>> a[b[x] for x in range(0,len(b))] = 1
  File "<stdin>", line 1
    a[b[x] for x in range(0,len(b))] = 1
             ^
SyntaxError: invalid syntax

Can anyone see what's going on here?

Thanks!


Solution

  • Using list comprehension

    In [1]: a=[0, 0, 0, 0]
    
    In [2]: b=[1, 3]
    
    In [3]: [ 1 if i in b else a_i for i, a_i in enumerate(a) ]
    Out[3]: [0, 1, 0, 1]
    

    "The elegant solution, if python was worth its salt, would be this"

    Import one module and python is, as you say, worth its salt:

    In [1]: from numpy import array
    
    In [2]: a = array([0, 0, 0, 0])
    
    In [3]: b = [1, 3]
    
    In [4]: a[b] = 1
    
    In [5]: a
    Out[5]: array([0, 1, 0, 1])
    

    For handling large quantities of data, numpy is both elegant and fast. If not using numpy, then Jonathan Hartnagel's for loop or something like Joel Cornett's BellsAndWhistlesList would be good solutions.