Python assignment in just one line

Question

I'm new to python and I saw this form to give values to a list

color= ['red' if v == 0 else 'green' for v in y]

But if I try to do it with 3 number, for example

color= ['red' if v == 0 elif v == 1 'blue' else 'green' for v in y]

is that possible or do I have to do it like this:

color = ['none']*len(y)
for i in range(0,len(color)):
 if y[i] == 0:
 color[i] = 'red'
 elif y[i] == 1:
 color[i] = 'blue'
 else:
 color[i] = 'green'

because this form is not as easy to write as the other one. Thank you.

• ababuji: Keep up your effort and remember you will always get help in SO if you've shown effort from your part which you have!

Answer 1

You can, but you have to modify your syntax slightly (and use else instead of the elifs):

color= ['red' if v == 0 else 'blue' if v == 1 else 'green' for v in y]

Example:

y = [1,0,2,3,0,1] 
color= ['red' if v == 0 else 'blue' if v == 1 else 'green' for v in y]
>>> color
['blue', 'red', 'green', 'green', 'red', 'blue']

• Jeffer Leonardo Alarcon Pedraz: Thank you, that's exactly what I want.

Answer 2

You could use a dictionary instead.

If you could restructure your questions with more details the answer would make more sense I think!

colors = {
 0: 'red',
 1: 'blue'
}
color = [colors.get(v, 'green') for v in y]

---

Dictionary class has a built-in function named .get() that accepts a key and a default value if the key is not found. colors.get(v, 'green') translates into: give me the value of key v in dictionary colors, however if not found, give me 'green'.

Answer 3

Since v has integer values, you can use the values to index into a list or tuple of the color names as:

>>> y = [0,1,2,1,1,0,2,2]
>>> [ ('red', 'blue')[i] if i in (0,1) else 'green' for i in y]
['red', 'blue', 'green', 'blue', 'blue', 'red', 'green', 'green']

Similar to the dictionary approaches, this approach could be extended to a longer list of items pretty easily, but you don't have to create a separate dictionary.

Answer 4

You could use a dictionary like so:

color_codes = { 0: 'red',
1: 'blue',
2: 'green'
}
color = [color_codes[v] for v in y]