CODEWITHSUNDEEP
pythonpython-3.xlist-comprehensionconditional-operator
fontesi
fontesi

Reputation: 15

Single-line if-elif-else with list comprehensions

item is all tuples from MySQL

y=[1 if (t['color']=='yellow') else -1 for t in item]

output:

[-1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1  -1 -1 -1 -1 -1 -1  1  1  1  1  1  1  1  1  1  1  1  1  1  1  1 ]

like this one, I want to do this

for t in item:
   if(t['color']=='yellow'):
       y=1
   elif(t['color']=='blue'):
       y=2
   else:
       y=3

and put the result in an array

y=[1 if (t['color']=='yellow') 2 elif (t['color']=='blue') else 3 for t in item]
y=np.array(y)
print(y)

output should be like this [ 1 2 3 1 3 2 1 3 1 2 3 2 1 ] How can I do that?

Upvotes: 1

Views: 85

Answers (2)

pho
pho

Reputation: 25489

This is a great use-case for a dictionary!

First, define a dict that maps the color names to the numbers

colordict = {'yellow': 1, 'blue': 2, 'default': 3}

Then, use colordict.get() to retrieve the correct number

y = [colordict.get(t['color'], colordict['default']) for t in item]

The second argument to .get() is the default value that it returns if t['color'] is not found in colordict.

Using a dict makes it trivial to add more colors. Can you imagine having to write a hundred nested if-elses to support a hundred colors?!

Testing with this dummy list:

item = [
    {'color': 'yellow', 'data': 0.5},
    {'color': 'purple', 'data': 0.1}, 
    {'color': 'blue', 'data': 0.2}, 
    {'color': 'blue', 'data': 0.3}, 
    {'color': 'red', 'data': 0.6}
       ]

we get the output

[1, 3, 2, 2, 3]

which is exactly what we expected.

You could also use defaultdict if you don't want to deal with .get().

Upvotes: 1

ysth
ysth

Reputation: 98398

y=[1 if (t['color']=='yellow') else 2 if (t['color']=='blue') else 3 for t in item]

Upvotes: 1

Related Questions