多个 if 语句列表理解

Question

Good evening,

I'm trying to get better at list comprehension since I discovered this fantastic form of coding a few days ago. Currently I'm trying to do a list comprehension with multiple if statements. This is the code I'm trying to rewrite (for your understanding)

for i in range(len(keys)):
    key = keys[i]

    if key == 1:
        newAns = [1, 0, 0, 0]
        answers.append(newAns)
    elif key == 2:
        newAns = [0, 1, 0, 0]
        answers.append(newAns)
    elif key == 3:
        newAns = [0, 0, 1, 0]
        answers.append(newAns)
    else:
        newAns = [0, 0, 0, 1]
        answers.append(newAns)

And this is what i have done so far

    answers = [i for i in keys]:
    [answers.append([1, 0, 0, 0]) if i == 1]
    [answers.append([0, 1, 0, 0]) if i == 2]
    [answers.append([0, 0, 1, 0]) if i == 3]
    [answers.append([0, 0, 0, 1]) if i == 1]

The list contains values of ints and i would like to convert them to vectors depending on what value the key has.

I'm a bit stuck and would appreciate some guidance in how to approach this task. Thank you.

Answer 1

How about we put all the key and newAns in a dict and use them in your list comprehension?

answer_map = {1: [1, 0, 0, 0], 2: [0, 1, 0, 0], 3: ...}
answers = [answer_map[x] if x in answer_map else [0, 0, 0, 1] for x in keys]

Update :

I totally forgot about dict.get(key, default) (Thanks, @U9-Forward!). You could also say:

[answer_map.get(x, [0, 0, 0, 1]) for x in keys]

Answer 2

Additional to @UltrInstinct's answer, do get to make no if statements:

answer_map = {1: [1, 0, 0, 0], 2: [0, 1, 0, 0], 3: ...}
answers = [answer_map,get(x,[0, 0, 0, 1]) for x in keys]

Now:

print(answers)

Would output the desired output.

Answer 3

As a fan of comprehensions your can use a nested comprehension one liner if

[ [ 1 if i == min(k, 4) else 0 for i in range(1, 5)]  for k in keys]

A two-liner with a dictionary

 answer_dict = {1: [1, 0, 0, 0],
               2: [0, 1, 0, 0] ...}

would be a bit more efficient. Just for fun you can construct the dictionary using a nested dictionary comprehension, even though it does not be worth the efort it if you have just four answers.

answer_dict = { k: [ 1 if i == k else 0
                   for i in range(1, 5) for k in range(1, 5)]
              }
[answer_dict[min(4, k)] for k in keys]   

Instead of dictionary you can also use a list

answers = [ [ 1 if i == k else 0
                   for i in range(1, 5) for k in range(0, 5)]]

[answers[min(4, k)] for k in keys] 

Answer 4

You could use nested ternary expressions if you want an absurd one-liner.

In [4]: keys = [2, 1, 1, 3, 0, 2]

In [5]: result  = [
   ...:     [1,0,0,0] if key == 1 else
   ...:     [0,1,0,0] if key == 2 else
   ...:     [0,0,1,0] if key == 3 else
   ...:     [0,0,0,1]
   ...:     for key in keys
   ...: ]

In [6]: result
Out[6]:
[[0, 1, 0, 0],
 [1, 0, 0, 0],
 [1, 0, 0, 0],
 [0, 0, 1, 0],
 [0, 0, 0, 1],
 [0, 1, 0, 0]]

However, it is better to wrap your logic in a function and call that function in the list comprehension:

In [7]: def f(key):
   ...:     if key == 1:
   ...:         result = [1, 0, 0, 0]
   ...:     elif key == 2:
   ...:         result = [0, 1, 0, 0]
   ...:     elif key == 3:
   ...:         result = [0, 0, 1, 0]
   ...:     else:
   ...:         result = [0, 0, 0, 1]
   ...:     return result
   ...:

In [8]: [f(key) for key in keys]
Out[8]:
[[0, 1, 0, 0],
 [1, 0, 0, 0],
 [1, 0, 0, 0],
 [0, 0, 1, 0],
 [0, 0, 0, 1],
 [0, 1, 0, 0]]

In this case, a dict works very nicely as well, as demonstrated in other answers. In general, don't try to cram a bunch of things into a single list-comprehension.