CODEWITHSUNDEEP
pythonlist-comprehensionfor-comprehension
Balthasar
Balthasar

Reputation: 1237

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!

Upvotes: 0

Views: 4081

Answers (8)

John1024
John1024

Reputation: 113814

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.

Upvotes: 7

Jonathan Hartnagel
Jonathan Hartnagel

Reputation: 193

What's wrong with :

for i in b:
    a[i] = 1

Or, as suggested by @John1024, using numpy :

from numpy import array
a[b] = 1

Upvotes: 3

Ejaz
Ejaz

Reputation: 1632

how about this:

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

for i in range(len(b)):
    a[b[i]] = 1

print a # [0, 1, 0, 1]

Upvotes: 0

Joel Cornett
Joel Cornett

Reputation: 24788

Just for kicks:

class BellsAndWhistlesList(list):
    def __setitem__(self, key, value):
        if hasattr(key, "__iter__") or hasattr(key, "next"):
            for i, v in zip(key, value):
                super(BellsAndWhistlesList, self).__setitem__(i, v)
        else:
            super(BellsAndWhistlesList, self).__setitem__(key, value)

Upvotes: 0

WKPlus
WKPlus

Reputation: 7255

I don't think [1 if i in b else v for i,v in enumerate(a)] is efficient, since it needs to check whether i is in b for each i in enumerate(a).

As other answers said, for i in b: a[i] = 1 is the simplest and efficient way, it just set values it needs to directly.

However, if you don't want to use the for loop, there is another way to use map instead of it to apply operations to all elements of a list. This implement will introduce an function call for each element, so maybe it's not efficient enough either.

>>> a=[0,0,0,0]
>>> b=[1,3]
>>> import functools
>>> def set_index(x, l):
...     l[x] = 1
...
>>> mapfunc = functools.partial(set_index, l=a)
>>> map(mapfunc, b)
[None, None]
>>> a
[0, 1, 0, 1]

Upvotes: 0

venpa
venpa

Reputation: 4318

Using list comprehension:

a = [1 if i in b else v for i,v in enumerate(a) ]

output:

[0, 1, 0, 1]

Upvotes: 1

Wolph
Wolph

Reputation: 80011

Try numpy if you want stuff like this:

In [1]: import numpy

In [2]: a = numpy.zeros(4)

In [3]: b = numpy.array((1, 3))

In [4]: a[b] = 1

In [5]: a
Out[5]: array([ 0.,  1.,  0.,  1.])

Or plain Python solutions like these:

a = [0] * 4
b = [1, 3]
for i in b:
    a[i] = 1

Alternatively (slow and inefficient):

a = [0] * 4
b = [1, 3]
a = [i in b and 1 or a_i for i, a_i in enumerate(a)]

Upvotes: 1

Maroun
Maroun

Reputation: 95948

[x for x in b]

Will create:

[1, 3]

So your code will look like:

a[1,3] = 1

Which is invalid.

I wouldn't use list comprehension for this, you can simply:

for i in b:
    a[i] = 1

Upvotes: 0

Related Questions